⚡ Support for model meta-llama-3.1-405b-instruct

Author: MrGraversen

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

@@ -5,6 +5,7 @@
 
 @UtilityClass
 public class Llama3Models {
+    public static final ReplicateModel LLAMA_3_1_405B_INSTRUCT = new ReplicateModel("meta", "meta-llama-3.1-405b-instruct");
     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/Llama3PredictionMapper.java

@@ -3,10 +3,7 @@
 import io.graversen.replicate.common.BasePredictionMapper;
 import io.graversen.replicate.common.PredictionTypes;
 import io.graversen.replicate.common.ReplicateModel;
-import io.graversen.replicate.models.MetaMetaLlama370BInstructInput;
-import io.graversen.replicate.models.MetaMetaLlama370BInstructPredictionrequest;
-import io.graversen.replicate.models.MetaMetaLlama38BInstructInput;
-import io.graversen.replicate.models.MetaMetaLlama38BInstructPredictionrequest;
+import io.graversen.replicate.models.*;
 import io.graversen.replicate.service.CreateTextPrediction;
 import lombok.NonNull;
 import org.springframework.stereotype.Component;
@@ -19,6 +16,7 @@ public class Llama3PredictionMapper extends BasePredictionMapper<CreateTextPredi
     @Override
     protected Set<ReplicateModel> supportedModels() {
         return Set.of(
+                Llama3Models.LLAMA_3_1_405B_INSTRUCT,
                 Llama3Models.LLAMA_3_70B_INSTRUCT,
                 Llama3Models.LLAMA_3_8B_INSTRUCT
         );
@@ -44,6 +42,15 @@ public Object apply(@NonNull ReplicateModel model, @NonNull CreateTextPrediction
                     null,
                     null
             );
+        } else if (model.equals(Llama3Models.LLAMA_3_1_405B_INSTRUCT)) {
+            return new MetaMetaLlama31405BInstructPredictionrequest(
+                    null,
+                    mapLlama31405bInput().apply(createPrediction),
+                    null,
+                    null,
+                    null,
+                    null
+            );
         } else {
             throw new IllegalArgumentException("Unsupported Replicate Model: " + model);
         }
@@ -54,6 +61,27 @@ public boolean supportsType(@NonNull PredictionTypes type) {
         return PredictionTypes.TEXT.equals(type);
     }
 
+    Function<CreateTextPrediction, MetaMetaLlama31405BInstructInput> mapLlama31405bInput() {
+        return createTextPrediction -> {
+            final var conversation = createTextPrediction.getConversation();
+            final var truncatedConversation = Llama3Tokenizer.fitToContextWindow(conversation, Llama3Tokenizer.DEFAULT_CONTEXT_WINDOW_SIZE);
+            final var textCompletion = Llama3Tokenizer.generateTextCompletion(truncatedConversation);
+
+            return new MetaMetaLlama31405BInstructInput(
+                    createTextPrediction.getTopK(),
+                    createTextPrediction.getTopP(),
+                    textCompletion.getText(),
+                    createTextPrediction.getMaxTokens(),
+                    createTextPrediction.getMinTokens(),
+                    createTextPrediction.getTemperature(),
+                    truncatedConversation.getSystemMessage(),
+                    null,
+                    null,
+                    null
+            );
+        };
+    }
+
     Function<CreateTextPrediction, MetaMetaLlama370BInstructInput> mapLlama370bInput() {
         return createTextPrediction -> {
             final var conversation = createTextPrediction.getConversation();