⚡ Generate TextConversation from tokenized conversation

Author: MrGraversen

spring-boot-starter-replicate/src/main/java/io/graversen/replicate/common/TextConversation.java

@@ -3,20 +3,38 @@
 import lombok.NonNull;
 import lombok.Value;
 
+import java.util.ArrayList;
 import java.util.List;
 
 @Value
 public class TextConversation {
     @NonNull String systemMessage;
     @NonNull List<TextMessage> messages;
 
+    public static TextConversation of(@NonNull String systemMessage) {
+        return new TextConversation(
+                systemMessage,
+                List.of()
+        );
+    }
+
     public static TextConversation of(@NonNull String systemMessage, @NonNull String userMessage) {
         return new TextConversation(
                 systemMessage,
                 List.of(TextMessage.user(userMessage))
         );
     }
 
+    public TextConversation append(@NonNull TextMessage message) {
+        final var mutableMessages = new ArrayList<>(getMessages());
+        mutableMessages.add(message);
+
+        return new TextConversation(
+                getSystemMessage(),
+                List.copyOf(mutableMessages)
+        );
+    }
+
     public List<TextMessage> getLastMessages(@NonNull Integer conversationSize) {
         return messages.stream()
                 .skip(Math.max(0, messages.size() - conversationSize))

spring-boot-starter-replicate/src/main/java/io/graversen/replicate/llama3/Llama3Tokenizer.java

@@ -6,9 +6,11 @@
 import lombok.NonNull;
 import lombok.experimental.UtilityClass;
 
+import java.util.ArrayList;
 import java.util.LinkedList;
 import java.util.Objects;
 import java.util.function.Consumer;
+import java.util.regex.Pattern;
 
 @UtilityClass
 public class Llama3Tokenizer {
@@ -60,6 +62,41 @@ public static Llama3TextCompletion generateTextCompletion(@NonNull TextConversat
         return new Llama3TextCompletion(textCompletion);
     }
 
+    public static TextConversation parseTextCompletion(@NonNull String tokenizedConversation) {
+        String systemMessage = null;
+        final var messages = new ArrayList<TextMessage>();
+
+        final String[] tokens = tokenizedConversation.split(Pattern.quote(END_OF_TEXT_ID));
+
+        for (String token : tokens) {
+            token = token.trim();
+
+            if (token.isEmpty()) {
+                continue;
+            }
+
+            if (token.startsWith(BEGIN_OF_TEXT)) {
+                String systemToken = token.substring(BEGIN_OF_TEXT.length()).trim();
+                if (systemToken.startsWith(START_HEADER_ID + ROLE_SYSTEM)) {
+                    systemMessage = systemToken.substring(systemToken.indexOf(END_HEADER_ID) + END_HEADER_ID.length()).trim();
+                }
+            } else if (token.startsWith(START_HEADER_ID)) {
+                String role = token.substring(START_HEADER_ID.length(), token.indexOf(END_HEADER_ID)).trim();
+                String messageContent = token.substring(token.indexOf(END_HEADER_ID) + END_HEADER_ID.length()).trim();
+
+                if (!role.isEmpty() && !messageContent.isEmpty()) {
+                    messages.add(new TextMessage(role, messageContent));
+                }
+            }
+        }
+
+        if (systemMessage == null) {
+            throw new IllegalArgumentException("System message not found in the tokenized conversation.");
+        }
+
+        return new TextConversation(systemMessage, messages);
+    }
+
     public static Integer approximateConversationContextSize(@NonNull TextConversation conversation, @Nullable Integer tokenSize) {
         final var conversationTextCompletion = generateTextCompletion(conversation);
         return getTokens(conversationTextCompletion.getText(), tokenSize);

spring-boot-starter-replicate/src/test/java/io/graversen/replicate/llama3/Llama3TokenizerTest.java

@@ -19,6 +19,38 @@ class Llama3TokenizerTest {
             new TextMessage("assistant", "Why don't scientists trust atoms? Because they make up everything!")
     );
 
+    @Test
+    void parseTextCompletion_exampleConversation() {
+        final var tokenizedConversation =
+                "<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n" +
+                "You are a helpful assistant.<|eot_id|><|start_header_id|>user<|end_header_id|>\n" +
+                "Hello. How are you?<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n" +
+                "Hello! I'm doing great, thanks for asking! I'm here to help you with anything you need, so please feel free to ask me any questions or share what's on your mind. How about you? How's your day going so far?<|eot_id|><|start_header_id|>user<|end_header_id|>\n" +
+                "It is going good thanks<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n" +
+                "That's wonderful to hear! I'm glad to know that your day is going well. If you don't mind me asking, what's been the highlight of your day so far? Is there anything exciting or interesting that's happened? I'm all ears and happy to listen!<|eot_id|><|start_header_id|>user<|end_header_id|>\n" +
+                "Nah fam<|eot_id|>";
+
+        final var conversation = Llama3Tokenizer.parseTextCompletion(tokenizedConversation);
+
+        assertEquals("You are a helpful assistant.", conversation.getSystemMessage());
+        assertEquals(5, conversation.getMessages().size());
+
+        assertEquals("user", conversation.getMessages().get(0).getRole());
+        assertEquals("Hello. How are you?", conversation.getMessages().get(0).getText());
+
+        assertEquals("assistant", conversation.getMessages().get(1).getRole());
+        assertEquals("Hello! I'm doing great, thanks for asking! I'm here to help you with anything you need, so please feel free to ask me any questions or share what's on your mind. How about you? How's your day going so far?", conversation.getMessages().get(1).getText());
+
+        assertEquals("user", conversation.getMessages().get(2).getRole());
+        assertEquals("It is going good thanks", conversation.getMessages().get(2).getText());
+
+        assertEquals("assistant", conversation.getMessages().get(3).getRole());
+        assertEquals("That's wonderful to hear! I'm glad to know that your day is going well. If you don't mind me asking, what's been the highlight of your day so far? Is there anything exciting or interesting that's happened? I'm all ears and happy to listen!", conversation.getMessages().get(3).getText());
+
+        assertEquals("user", conversation.getMessages().get(4).getRole());
+        assertEquals("Nah fam", conversation.getMessages().get(4).getText());
+    }
+
     @Test
     public void fitToContextWindow_defaultWindow() {
         final var conversation = new TextConversation(systemMessage, messages);