feat: Support Llama 3 model (#479)

reneleonhardt · web-flow · commit 6e6a4991055d · 2024-04-21T01:12:13.000+03:00
* feat: Support Llama 3 model (#478) * Use new InfillPrompt * Switch to lmstudio-community * Use new Prompt * llama.cpp removed the BOS token ggml-org/llama.cpp@a55d8a9 * Add tests * I would prefer a stream based solution * Add 70B models * Add tests for skipping blank system prompt * Remove InfillPrompt for now
diff --git a/src/main/java/ee/carlrobert/codegpt/completions/HuggingFaceModel.java b/src/main/java/ee/carlrobert/codegpt/completions/HuggingFaceModel.java
@@ -43,16 +43,31 @@ public enum HuggingFaceModel {
   WIZARD_CODER_PYTHON_13B_Q5(13, 5, "WizardCoder-Python-13B-V1.0-GGUF"),
   WIZARD_CODER_PYTHON_34B_Q3(34, 3, "WizardCoder-Python-34B-V1.0-GGUF"),
   WIZARD_CODER_PYTHON_34B_Q4(34, 4, "WizardCoder-Python-34B-V1.0-GGUF"),
-  WIZARD_CODER_PYTHON_34B_Q5(34, 5, "WizardCoder-Python-34B-V1.0-GGUF");
+  WIZARD_CODER_PYTHON_34B_Q5(34, 5, "WizardCoder-Python-34B-V1.0-GGUF"),
+
+  LLAMA_3_8B_IQ3_M(8, 3, "Meta-Llama-3-8B-Instruct-IQ3_M.gguf", "lmstudio-community"),
+  LLAMA_3_8B_Q4_K_M(8, 4, "Meta-Llama-3-8B-Instruct-Q4_K_M.gguf", "lmstudio-community"),
+  LLAMA_3_8B_Q5_K_M(8, 5, "Meta-Llama-3-8B-Instruct-Q5_K_M.gguf", "lmstudio-community"),
+  LLAMA_3_8B_Q6_K(8, 6, "Meta-Llama-3-8B-Instruct-Q6_K.gguf", "lmstudio-community"),
+  LLAMA_3_8B_Q8_0(8, 8, "Meta-Llama-3-8B-Instruct-Q8_0.gguf", "lmstudio-community"),
+  LLAMA_3_70B_IQ1(70, 1, "Meta-Llama-3-70B-Instruct-IQ1_M.gguf", "lmstudio-community"),
+  LLAMA_3_70B_IQ2_XS(70, 2, "Meta-Llama-3-70B-Instruct-IQ2_XS.gguf", "lmstudio-community"),
+  LLAMA_3_70B_Q4_K_M(70, 4, "Meta-Llama-3-70B-Instruct-Q4_K_M.gguf", "lmstudio-community");
 
   private final int parameterSize;
   private final int quantization;
   private final String modelName;
+  private final String user;
 
   HuggingFaceModel(int parameterSize, int quantization, String modelName) {
+    this(parameterSize, quantization, modelName, "TheBloke");
+  }
+
+  HuggingFaceModel(int parameterSize, int quantization, String modelName, String user) {
     this.parameterSize = parameterSize;
     this.quantization = quantization;
     this.modelName = modelName;
+    this.user = user;
   }
 
   public int getParameterSize() {
@@ -68,26 +83,37 @@ public String getCode() {
   }
 
   public String getFileName() {
-    return modelName.toLowerCase().replace("-gguf", format(".Q%d_K_M.gguf", quantization));
+    if ("TheBloke".equals(user)) {
+      return modelName.toLowerCase().replace("-gguf", format(".Q%d_K_M.gguf", quantization));
+    }
+    return modelName;
   }
 
   public URL getFileURL() {
     try {
       return new URL(
-          format("https://huggingface.co/TheBloke/%s/resolve/main/%s", modelName, getFileName()));
+          "https://huggingface.co/%s/%s/resolve/main/%s".formatted(user, getDirectory(), getFileName()));
     } catch (MalformedURLException ex) {
       throw new RuntimeException(ex);
     }
   }
 
   public URL getHuggingFaceURL() {
     try {
-      return new URL("https://huggingface.co/TheBloke/" + modelName);
+      return new URL("https://huggingface.co/%s/%s".formatted(user, getDirectory()));
     } catch (MalformedURLException ex) {
       throw new RuntimeException(ex);
     }
   }
 
+  private String getDirectory() {
+    if ("lmstudio-community".equals(user)) {
+      // Meta-Llama-3-8B-Instruct-Q4_K_M.gguf -> Meta-Llama-3-8B-Instruct-GGUF
+      return modelName.replaceFirst("-[^.-]+\\.gguf$", "-GGUF");
+    }
+    return modelName;
+  }
+
   @Override
   public String toString() {
     return format("%d-bit precision", quantization);
diff --git a/src/main/java/ee/carlrobert/codegpt/completions/llama/LlamaModel.java b/src/main/java/ee/carlrobert/codegpt/completions/llama/LlamaModel.java
@@ -82,7 +82,24 @@ public enum LlamaModel {
           HuggingFaceModel.WIZARD_CODER_PYTHON_13B_Q5,
           HuggingFaceModel.WIZARD_CODER_PYTHON_34B_Q3,
           HuggingFaceModel.WIZARD_CODER_PYTHON_34B_Q4,
-          HuggingFaceModel.WIZARD_CODER_PYTHON_34B_Q5));
+          HuggingFaceModel.WIZARD_CODER_PYTHON_34B_Q5)),
+  LLAMA_3(
+      "Llama 3",
+      "Llama 3 is a family of large language models (LLMs), a collection of pretrained and "
+          + "instruction tuned generative text models in 8 and 70B sizes. The Llama 3 instruction "
+          + "tuned models are optimized for dialogue use cases and outperform many of the available"
+          + " open source chat models on common industry benchmarks. Further, in developing these "
+          + "models, we took great care to optimize helpfulness and safety.",
+      PromptTemplate.LLAMA_3,
+      List.of(
+          HuggingFaceModel.LLAMA_3_8B_IQ3_M,
+          HuggingFaceModel.LLAMA_3_8B_Q4_K_M,
+          HuggingFaceModel.LLAMA_3_8B_Q5_K_M,
+          HuggingFaceModel.LLAMA_3_8B_Q6_K,
+          HuggingFaceModel.LLAMA_3_8B_Q8_0,
+          HuggingFaceModel.LLAMA_3_70B_IQ1,
+          HuggingFaceModel.LLAMA_3_70B_IQ2_XS,
+          HuggingFaceModel.LLAMA_3_70B_Q4_K_M));
 
   private final String label;
   private final String description;
diff --git a/src/main/java/ee/carlrobert/codegpt/completions/llama/PromptTemplate.java b/src/main/java/ee/carlrobert/codegpt/completions/llama/PromptTemplate.java
@@ -1,7 +1,11 @@
 package ee.carlrobert.codegpt.completions.llama;
 
+import static java.util.stream.Stream.concat;
+
 import ee.carlrobert.codegpt.conversations.message.Message;
 import java.util.List;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 public enum PromptTemplate {
 
@@ -55,6 +59,26 @@ public String buildPrompt(String systemPrompt, String userPrompt, List<Message>
           .toString();
     }
   },
+  LLAMA_3("Llama 3") {
+    @Override
+    public String buildPrompt(String systemPrompt, String userPrompt, List<Message> history) {
+      return concat(concat(Stream.ofNullable(systemPrompt)
+                      .filter(s -> !s.isBlank())
+                      .flatMap(system -> Stream.of(
+                              "<|start_header_id|>system<|end_header_id|>\n\n",
+                              system,
+                              "<|eot_id|>")),
+              history.stream().flatMap(message -> mapMessage(
+                      message,
+                      "<|start_header_id|>user<|end_header_id|>\n\n",
+                      "<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n",
+                      "<|eot_id|>"))), Stream.of(
+              "<|start_header_id|>user<|end_header_id|>\n\n",
+              userPrompt,
+              "<|eot_id|>"))
+              .collect(Collectors.joining());
+    }
+  },
   MIXTRAL_INSTRUCT("Mixtral Instruct") {
     @Override
     public String buildPrompt(String systemPrompt, String userPrompt, List<Message> history) {
@@ -171,4 +195,15 @@ public String buildPrompt(String systemPrompt, String userPrompt, List<Message>
   public String toString() {
     return label;
   }
+
+  private static Stream<String> mapMessage(Message message,
+                                           String prefix, String infix, String suffix) {
+    return Stream.of(
+            prefix,
+            message.getPrompt(),
+            infix,
+            message.getResponse(),
+            suffix
+    );
+  }
 }
diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillPromptTemplate.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillPromptTemplate.kt
@@ -28,4 +28,4 @@ enum class InfillPromptTemplate(val label: String, val stopTokens: List<String>?
     override fun toString(): String {
         return label
     }
-}
+}
diff --git a/src/test/kotlin/ee/carlrobert/codegpt/completions/PromptTemplateTest.kt b/src/test/kotlin/ee/carlrobert/codegpt/completions/PromptTemplateTest.kt
@@ -3,10 +3,14 @@ package ee.carlrobert.codegpt.completions
 import ee.carlrobert.codegpt.completions.llama.PromptTemplate.ALPACA
 import ee.carlrobert.codegpt.completions.llama.PromptTemplate.CHAT_ML
 import ee.carlrobert.codegpt.completions.llama.PromptTemplate.LLAMA
+import ee.carlrobert.codegpt.completions.llama.PromptTemplate.LLAMA_3
 import ee.carlrobert.codegpt.completions.llama.PromptTemplate.TORA
 import ee.carlrobert.codegpt.conversations.message.Message
 import org.assertj.core.api.Assertions.assertThat
-import org.junit.Test
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.params.ParameterizedTest
+import org.junit.jupiter.params.provider.NullAndEmptySource
+import org.junit.jupiter.params.provider.ValueSource
 
 class PromptTemplateTest {
 
@@ -34,6 +38,72 @@ class PromptTemplateTest {
             """.trimIndent())
   }
 
+  @Test
+  fun shouldBuildLlama3PromptWithoutHistory() {
+    val prompt = LLAMA_3.buildPrompt(SYSTEM_PROMPT, USER_PROMPT, listOf())
+
+    assertThat(prompt).isEqualTo("""
+      <|start_header_id|>system<|end_header_id|>
+
+      TEST_SYSTEM_PROMPT<|eot_id|><|start_header_id|>user<|end_header_id|>
+
+      TEST_USER_PROMPT<|eot_id|>""".trimIndent()
+    )
+  }
+
+  @ParameterizedTest
+  @NullAndEmptySource
+  @ValueSource(strings = [" ", "\t", "\n"])
+  fun shouldBuildLlama3PromptWithoutHistorySkippingBlankSystemPrompt(systemPrompt: String?) {
+    val prompt = LLAMA_3.buildPrompt(systemPrompt, USER_PROMPT, listOf())
+
+    assertThat(prompt).isEqualTo("""
+      <|start_header_id|>user<|end_header_id|>
+
+      TEST_USER_PROMPT<|eot_id|>""".trimIndent()
+    )
+  }
+
+  @Test
+  fun shouldBuildLlama3PromptWithHistory() {
+    val prompt = LLAMA_3.buildPrompt(SYSTEM_PROMPT, USER_PROMPT, HISTORY)
+
+    assertThat(prompt).isEqualTo("""
+      <|start_header_id|>system<|end_header_id|>
+
+      TEST_SYSTEM_PROMPT<|eot_id|><|start_header_id|>user<|end_header_id|>
+  
+      TEST_PREV_PROMPT_1<|eot_id|><|start_header_id|>assistant<|end_header_id|>
+
+      TEST_PREV_RESPONSE_1<|eot_id|><|start_header_id|>user<|end_header_id|>
+
+      TEST_PREV_PROMPT_2<|eot_id|><|start_header_id|>assistant<|end_header_id|>
+
+      TEST_PREV_RESPONSE_2<|eot_id|><|start_header_id|>user<|end_header_id|>
+
+      TEST_USER_PROMPT<|eot_id|>""".trimIndent())
+  }
+
+  @ParameterizedTest
+  @NullAndEmptySource
+  @ValueSource(strings = [" ", "\t", "\n"])
+  fun shouldBuildLlama3PromptWithHistorySkippingBlankSystemPrompt(systemPrompt: String?) {
+    val prompt = LLAMA_3.buildPrompt(systemPrompt, USER_PROMPT, HISTORY)
+
+    assertThat(prompt).isEqualTo("""
+      <|start_header_id|>user<|end_header_id|>
+  
+      TEST_PREV_PROMPT_1<|eot_id|><|start_header_id|>assistant<|end_header_id|>
+
+      TEST_PREV_RESPONSE_1<|eot_id|><|start_header_id|>user<|end_header_id|>
+
+      TEST_PREV_PROMPT_2<|eot_id|><|start_header_id|>assistant<|end_header_id|>
+
+      TEST_PREV_RESPONSE_2<|eot_id|><|start_header_id|>user<|end_header_id|>
+
+      TEST_USER_PROMPT<|eot_id|>""".trimIndent())
+  }
+
   @Test
   fun shouldBuildAlpacaPromptWithHistory() {
     val prompt = ALPACA.buildPrompt(SYSTEM_PROMPT, USER_PROMPT, HISTORY)

Original file line number	Diff line number	Diff line change
`@@ -28,4 +28,4 @@ enum class InfillPromptTemplate(val label: String, val stopTokens: List<String>?`
`28`	`28`	`override fun toString(): String {`
`29`	`29`	`return label`
`30`	`30`	`}`
`31`		`-}`
	`31`	`+}`