⚡ Llama3Tokenizer, context window conversation fitting

Author: MrGraversen

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

@@ -0,0 +1,107 @@
+package io.graversen.replicate.llama3;
+
+import io.graversen.replicate.common.TextConversation;
+import io.graversen.replicate.common.TextMessage;
+import jakarta.annotation.Nullable;
+import lombok.NonNull;
+import lombok.experimental.UtilityClass;
+
+import java.util.LinkedList;
+import java.util.Objects;
+import java.util.function.Consumer;
+import java.util.stream.Collectors;
+
+@UtilityClass
+public class Llama3Tokenizer {
+    private static final String BEGIN_OF_TEXT = "<|begin_of_text|>";
+    private static final String START_HEADER_ID = "<|start_header_id|>";
+    private static final String END_HEADER_ID = "<|end_header_id|>";
+    private static final String END_OF_TEXT_ID = "<|eot_id|>";
+
+    public static final Integer DEFAULT_CONTEXT_WINDOW_SIZE = 8000;
+    public static final Integer APPROXIMATE_CHARACTERS_PER_TOKEN = 4;
+    public static final String ROLE_USER = "user";
+    public static final String ROLE_ASSISTANT = "assistant";
+    public static final String ROLE_SYSTEM = "system";
+
+    public static String beginOfText(@NonNull String text) {
+        return String.format("%s%s", BEGIN_OF_TEXT, text);
+    }
+
+    public static String endOfText(@NonNull String text) {
+        return String.format("%s%s", text, END_OF_TEXT_ID);
+    }
+
+    public static String header(@NonNull String text) {
+        return String.format("%s%s%s", START_HEADER_ID, text, END_HEADER_ID);
+    }
+
+    public static String userHeader() {
+        return Llama3Tokenizer.header(ROLE_USER);
+    }
+
+    public static String assistantHeader() {
+        return Llama3Tokenizer.header(ROLE_ASSISTANT);
+    }
+
+    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());
+
+        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);
+    }
+
+    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);
+        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);
+
+            if (remainingTokens - messageTokens >= 0) {
+                fittedMessages.addFirst(message);
+                remainingTokens -= messageTokens;
+            } else {
+                break;
+            }
+        }
+
+        return new TextConversation(systemMessage, fittedMessages);
+    }
+
+    Consumer<TextMessage> addMessageToTextCompletion(@NonNull StringBuilder textCompletionBuilder) {
+        return textMessage -> textCompletionBuilder
+                .append(header(textMessage.getRole()))
+                .append(textMessage.getText())
+                .append(END_OF_TEXT_ID);
+    }
+
+    private Integer getTokens(@NonNull String string, @Nullable Integer tokenSize) {
+        tokenSize = Objects.requireNonNullElse(tokenSize, APPROXIMATE_CHARACTERS_PER_TOKEN);
+        return string.length() / tokenSize;
+    }
+}

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

@@ -0,0 +1,47 @@
+package io.graversen.replicate.llama3;
+
+import io.graversen.replicate.common.TextConversation;
+import io.graversen.replicate.common.TextMessage;
+import org.junit.jupiter.api.Test;
+
+import java.util.List;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+class Llama3TokenizerTest {
+    private final String systemMessage = "You are a helpful assistant";
+    private final List<TextMessage> messages = List.of(
+            new TextMessage("user", "Hello"),
+            new TextMessage("assistant", "Hi there!"),
+            new TextMessage("user", "How are you?"),
+            new TextMessage("assistant", "I'm fine, thank you! How can I assist you today?"),
+            new TextMessage("user", "Tell me a joke."),
+            new TextMessage("assistant", "Why don't scientists trust atoms? Because they make up everything!")
+    );
+
+    @Test
+    public void fitToContextWindow_defaultWindow() {
+        final var conversation = new TextConversation(systemMessage, messages);
+        final var fittedConversation = Llama3Tokenizer.fitToContextWindow(conversation, null);
+
+        assertNotNull(fittedConversation);
+        assertTrue(fittedConversation.getMessages().size() <= messages.size());
+        assertEquals(systemMessage, fittedConversation.getSystemMessage());
+        assertEquals(messages.get(0), fittedConversation.getMessages().get(0));
+        assertEquals(messages.get(5), fittedConversation.getMessages().get(5));
+        fittedConversation.getMessages().forEach(message ->
+                assertTrue(messages.contains(message), "Fitted conversation should only contain messages from the original conversation")
+        );
+    }
+
+    @Test
+    public void fitToContextWindow_windowTooSmall() {
+        final var conversation = new TextConversation(systemMessage, messages);
+        final var fittedConversation = Llama3Tokenizer.fitToContextWindow(conversation, 32);
+
+        assertNotNull(fittedConversation);
+        assertTrue(fittedConversation.getMessages().size() < messages.size());
+        assertNotEquals(messages.get(0), fittedConversation.getMessages().get(0));
+        assertEquals(messages.get(5), fittedConversation.getMessages().get(1));
+    }
+}