⚡ Refine Llama3 token approximation

Author: MrGraversen

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

@@ -13,9 +13,19 @@ public class TextConversation {
     public static TextConversation of(@NonNull String systemMessage, @NonNull String userMessage) {
         return new TextConversation(
                 systemMessage,
-                List.of(
-                        TextMessage.user(userMessage)
-                )
+                List.of(TextMessage.user(userMessage))
         );
     }
+
+    public List<TextMessage> getLastMessages(@NonNull Integer conversationSize) {
+        return messages.stream()
+                .skip(Math.max(0, messages.size() - conversationSize))
+                .toList();
+    }
+
+    public List<TextMessage> getFirstMessages(@NonNull Integer conversationSize) {
+        return messages.stream()
+                .limit(conversationSize)
+                .toList();
+    }
 }

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

@@ -9,7 +9,6 @@
 import java.util.LinkedList;
 import java.util.Objects;
 import java.util.function.Consumer;
-import java.util.stream.Collectors;
 
 @UtilityClass
 public class Llama3Tokenizer {
@@ -48,39 +47,38 @@ public static String systemHeader() {
         return Llama3Tokenizer.header(ROLE_SYSTEM);
     }
 
-    public static Llama3TextCompletion generateTextCompletion(@NonNull TextConversation conversation) {
-        final var textCompletionBuilder = new StringBuilder();
-        textCompletionBuilder
-                .append(BEGIN_OF_TEXT)
-                .append(systemHeader())
-                .append(conversation.getSystemMessage());
+    public static String tokenizeMessage(@NonNull TextMessage message) {
+        final var messageBuilder = new StringBuilder();
+        addMessageToTextCompletion(messageBuilder).accept(message);
+        return messageBuilder.toString();
+    }
 
+    public static Llama3TextCompletion generateTextCompletion(@NonNull TextConversation conversation) {
+        final var textCompletionBuilder = createBeginOfText(conversation.getSystemMessage());
         conversation.getMessages().forEach(addMessageToTextCompletion(textCompletionBuilder));
         final var textCompletion = textCompletionBuilder.toString();
         return new Llama3TextCompletion(textCompletion);
     }
 
     public static Integer approximateConversationContextSize(@NonNull TextConversation conversation, @Nullable Integer tokenSize) {
-        final var conversationPlainText = conversation.getMessages().stream()
-                .map(TextMessage::getText)
-                .collect(Collectors.joining(System.lineSeparator()));
-
-        return getTokens(conversationPlainText, tokenSize);
+        final var conversationTextCompletion = generateTextCompletion(conversation);
+        return getTokens(conversationTextCompletion.getText(), tokenSize);
     }
 
     public static TextConversation fitToContextWindow(@NonNull TextConversation conversation, @Nullable Integer contextWindowSize) {
         contextWindowSize = Objects.requireNonNullElse(contextWindowSize, DEFAULT_CONTEXT_WINDOW_SIZE);
 
-        final var systemMessage = conversation.getSystemMessage();
-        final var systemMessageTokens = getTokens(systemMessage, null);
+        final var systemMessage = createBeginOfText(conversation.getSystemMessage());
+        final var systemMessageTokens = getTokens(systemMessage.toString(), null);
         int remainingTokens = contextWindowSize - systemMessageTokens;
 
         final var messages = conversation.getMessages();
         final var fittedMessages = new LinkedList<TextMessage>();
 
         for (int i = messages.size() - 1; i >= 0; i--) {
             final var message = messages.get(i);
-            final var messageTokens = getTokens(message.getText(), null);
+            final var tokenizedMessage = tokenizeMessage(message);
+            final var messageTokens = getTokens(tokenizedMessage, null);
 
             if (remainingTokens - messageTokens >= 0) {
                 fittedMessages.addFirst(message);
@@ -90,7 +88,7 @@ public static TextConversation fitToContextWindow(@NonNull TextConversation conv
             }
         }
 
-        return new TextConversation(systemMessage, fittedMessages);
+        return new TextConversation(conversation.getSystemMessage(), fittedMessages);
     }
 
     Consumer<TextMessage> addMessageToTextCompletion(@NonNull StringBuilder textCompletionBuilder) {
@@ -100,8 +98,14 @@ Consumer<TextMessage> addMessageToTextCompletion(@NonNull StringBuilder textComp
                 .append(END_OF_TEXT_ID);
     }
 
-    private Integer getTokens(@NonNull String string, @Nullable Integer tokenSize) {
+    Integer getTokens(@NonNull String string, @Nullable Integer tokenSize) {
         tokenSize = Objects.requireNonNullElse(tokenSize, APPROXIMATE_CHARACTERS_PER_TOKEN);
         return string.length() / tokenSize;
     }
+
+    StringBuilder createBeginOfText(@NonNull String systemMessage) {
+        return new StringBuilder().append(BEGIN_OF_TEXT)
+                .append(systemHeader())
+                .append(systemMessage);
+    }
 }

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

@@ -37,11 +37,11 @@ public void fitToContextWindow_defaultWindow() {
     @Test
     public void fitToContextWindow_windowTooSmall() {
         final var conversation = new TextConversation(systemMessage, messages);
-        final var fittedConversation = Llama3Tokenizer.fitToContextWindow(conversation, 32);
+        final var fittedConversation = Llama3Tokenizer.fitToContextWindow(conversation, 128);
 
         assertNotNull(fittedConversation);
         assertTrue(fittedConversation.getMessages().size() < messages.size());
         assertNotEquals(messages.get(0), fittedConversation.getMessages().get(0));
-        assertEquals(messages.get(5), fittedConversation.getMessages().get(1));
+        assertEquals(messages.get(2), fittedConversation.getMessages().get(1));
     }
 }