Skip to content

Commit 7e69d9e

Browse files
authored
Merge pull request #463 from google-ai-edge/mrschmidt/spaces
Trim spaces in UI display
2 parents bdafe5d + 8b6ec73 commit 7e69d9e

File tree

2 files changed

+13
-11
lines changed

2 files changed

+13
-11
lines changed

examples/llm_inference/android/app/src/main/java/com/google/mediapipe/examples/llminference/ChatMessage.kt

+3-1
Original file line numberDiff line numberDiff line change
@@ -7,10 +7,12 @@ import java.util.UUID
77
*/
88
data class ChatMessage(
99
val id: String = UUID.randomUUID().toString(),
10-
val message: String = "",
10+
val rawMessage: String = "",
1111
val author: String,
1212
val isLoading: Boolean = false
1313
) {
1414
val isFromUser: Boolean
1515
get() = author == USER_PREFIX
16+
val message: String
17+
get() = rawMessage.trim()
1618
}

examples/llm_inference/android/app/src/main/java/com/google/mediapipe/examples/llminference/ChatUiState.kt

+10-10
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@ class ChatUiState(
3939

4040
// Prompt the model with the current chat history
4141
override val fullPrompt: String
42-
get() = _messages.joinToString(separator = "\n") { it.message }
42+
get() = _messages.joinToString(separator = "\n") { it.rawMessage }
4343

4444
override fun createLoadingMessage(): String {
4545
val chatMessage = ChatMessage(author = MODEL_PREFIX, isLoading = true)
@@ -54,14 +54,14 @@ class ChatUiState(
5454
override fun appendMessage(id: String, text: String, done: Boolean) {
5555
val index = _messages.indexOfFirst { it.id == id }
5656
if (index != -1) {
57-
val newText = _messages[index].message + text
58-
_messages[index] = _messages[index].copy(message = newText, isLoading = false)
57+
val newText = _messages[index].rawMessage + text
58+
_messages[index] = _messages[index].copy(rawMessage = newText, isLoading = false)
5959
}
6060
}
6161

6262
override fun addMessage(text: String, author: String): String {
6363
val chatMessage = ChatMessage(
64-
message = text,
64+
rawMessage = text,
6565
author = author
6666
)
6767
_messages.add(chatMessage)
@@ -85,7 +85,7 @@ class GemmaUiState(
8585
_messages. apply{
8686
for (i in indices) {
8787
this[i] = this[i].copy(
88-
message = this[i].message.replace(START_TURN + this[i].author + "\n", "")
88+
rawMessage = this[i].rawMessage.replace(START_TURN + this[i].author + "\n", "")
8989
.replace(END_TURN, "")
9090
)
9191
}
@@ -95,7 +95,7 @@ class GemmaUiState(
9595

9696
// Only using the last 4 messages to keep input + output short
9797
override val fullPrompt: String
98-
get() = _messages.takeLast(4).joinToString(separator = "\n") { it.message }
98+
get() = _messages.takeLast(4).joinToString(separator = "\n") { it.rawMessage }
9999

100100
override fun createLoadingMessage(): String {
101101
val chatMessage = ChatMessage(author = MODEL_PREFIX, isLoading = true)
@@ -112,18 +112,18 @@ class GemmaUiState(
112112
if (index != -1) {
113113
val newText = if (done) {
114114
// Append the Suffix when model is done generating the response
115-
_messages[index].message + text + END_TURN
115+
_messages[index].rawMessage + text + END_TURN
116116
} else {
117117
// Append the text
118-
_messages[index].message + text
118+
_messages[index].rawMessage + text
119119
}
120-
_messages[index] = _messages[index].copy(message = newText, isLoading = false)
120+
_messages[index] = _messages[index].copy(rawMessage = newText, isLoading = false)
121121
}
122122
}
123123

124124
override fun addMessage(text: String, author: String): String {
125125
val chatMessage = ChatMessage(
126-
message = "$START_TURN$author\n$text$END_TURN",
126+
rawMessage = "$START_TURN$author\n$text$END_TURN",
127127
author = author
128128
)
129129
_messages.add(chatMessage)

0 commit comments

Comments
 (0)