⚗️ Generate Llama3 conversation using LLM tokens

Author: MrGraversen

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

@@ -0,0 +1,21 @@
+package io.graversen.replicate.common;
+
+import lombok.NonNull;
+import lombok.Value;
+
+import java.util.List;
+
+@Value
+public class TextConversation {
+    @NonNull String systemMessage;
+    @NonNull List<TextMessage> messages;
+
+    public static TextConversation of(@NonNull String systemMessage, @NonNull String userMessage) {
+        return new TextConversation(
+                systemMessage,
+                List.of(
+                        TextMessage.user(userMessage)
+                )
+        );
+    }
+}

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

@@ -0,0 +1,18 @@
+package io.graversen.replicate.common;
+
+import lombok.NonNull;
+import lombok.Value;
+
+@Value
+public class TextMessage {
+    @NonNull String role;
+    @NonNull String text;
+
+    public static TextMessage user(@NonNull String text) {
+        return new TextMessage(TextPredictionRoles.USER.asString(), text);
+    }
+
+    public static TextMessage assistant(@NonNull String text) {
+        return new TextMessage(TextPredictionRoles.ASSISTANT.asString(), text);
+    }
+}

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

@@ -0,0 +1,10 @@
+package io.graversen.replicate.common;
+
+public enum TextPredictionRoles {
+    USER,
+    ASSISTANT;
+
+    public String asString() {
+        return name().toLowerCase();
+    }
+}

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

@@ -0,0 +1,10 @@
+package io.graversen.replicate.llama3;
+
+import io.graversen.replicate.common.ReplicateModel;
+import lombok.experimental.UtilityClass;
+
+@UtilityClass
+public class Llama3Models {
+    public static final ReplicateModel LLAMA_3_70B_INSTRUCT = new ReplicateModel("meta", "meta-llama-3-70b-instruct");
+    public static final ReplicateModel LLAMA_3_8B_INSTRUCT = new ReplicateModel("meta", "meta-llama-3-8b-instruct");
+}

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

@@ -0,0 +1,11 @@
+package io.graversen.replicate.llama3;
+
+import lombok.Getter;
+import lombok.NonNull;
+import lombok.RequiredArgsConstructor;
+
+@Getter
+@RequiredArgsConstructor
+public class Llama3TextCompletion {
+    private final @NonNull String text;
+}

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

@@ -0,0 +1,63 @@
+package io.graversen.replicate.llama3;
+
+import io.graversen.replicate.common.TextConversation;
+import io.graversen.replicate.common.TextMessage;
+import lombok.NonNull;
+import lombok.experimental.UtilityClass;
+
+import java.util.function.Consumer;
+
+@UtilityClass
+public class Llama3Tokens {
+    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 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 Llama3Tokens.header(ROLE_USER);
+    }
+
+    public static String assistantHeader() {
+        return Llama3Tokens.header(ROLE_ASSISTANT);
+    }
+
+    public static String systemHeader() {
+        return Llama3Tokens.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);
+    }
+
+    Consumer<TextMessage> addMessageToTextCompletion(@NonNull StringBuilder textCompletionBuilder) {
+        return textMessage -> textCompletionBuilder
+                .append(header(textMessage.getRole()))
+                .append(textMessage.getText())
+                .append(END_OF_TEXT_ID);
+    }
+}