feat: add CodeGemma InfillPromptTemplate (#530)

This commit is contained in:
Phil 2024-05-07 16:51:04 +02:00 committed by GitHub
parent a2a8747aca
commit 2c0a28a912
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 17 additions and 8 deletions

View file

@ -17,7 +17,7 @@ public enum LlamaModel {
+ "support for large input contexts, and zero-shot instruction following ability for "
+ "programming tasks.",
PromptTemplate.LLAMA,
InfillPromptTemplate.LLAMA,
InfillPromptTemplate.CODE_LLAMA,
List.of(
HuggingFaceModel.CODE_LLAMA_7B_Q3,
HuggingFaceModel.CODE_LLAMA_7B_Q4,
@ -115,12 +115,13 @@ public enum LlamaModel {
HuggingFaceModel.PHI_3_3_8B_4K_Q8_0,
HuggingFaceModel.PHI_3_3_8B_4K_FP16)),
CODE_GEMMA(
"CodeGemma 7b Instruct",
"CodeGemma 7b Instruct is the first in a series of coding models released by Google. "
"CodeGemma Instruct",
"CodeGemma Instruct is the first in a series of coding models released by Google. "
+ "As an instruct model, it specializes in being asked coding related questions, but can "
+ "also function as an autocomplete/fill-in-middle model for tools like co-pilot.\n"
+ "This model is perfect for general coding questions or code generation.",
PromptTemplate.CODE_GEMMA,
InfillPromptTemplate.CODE_GEMMA,
List.of(
HuggingFaceModel.CODE_GEMMA_7B_IQ1_S,
HuggingFaceModel.CODE_GEMMA_7B_IQ1_M,

View file

@ -144,7 +144,7 @@ public enum PromptTemplate {
.toString();
}
},
CODE_GEMMA("CodeGemma 7b Instruct") {
CODE_GEMMA("CodeGemma Instruct") {
@Override
public String buildPrompt(String systemPrompt, String userPrompt, List<Message> history) {
StringBuilder prompt = new StringBuilder();

View file

@ -16,8 +16,8 @@ public class LlamaSettingsState {
private HuggingFaceModel huggingFaceModel = HuggingFaceModel.CODE_LLAMA_7B_Q4;
private PromptTemplate localModelPromptTemplate = PromptTemplate.LLAMA;
private PromptTemplate remoteModelPromptTemplate = PromptTemplate.LLAMA;
private InfillPromptTemplate localModelInfillPromptTemplate = InfillPromptTemplate.LLAMA;
private InfillPromptTemplate remoteModelInfillPromptTemplate = InfillPromptTemplate.LLAMA;
private InfillPromptTemplate localModelInfillPromptTemplate = InfillPromptTemplate.CODE_LLAMA;
private InfillPromptTemplate remoteModelInfillPromptTemplate = InfillPromptTemplate.CODE_LLAMA;
private String baseHost = "http://localhost:8080";
private Integer serverPort = getRandomAvailablePortOrDefault();
private int contextSize = 2048;

View file

@ -7,11 +7,19 @@ enum class InfillPromptTemplate(val label: String, val stopTokens: List<String>?
return "<|fim_prefix|> $prefix <|fim_suffix|>$suffix <|fim_middle|>"
}
},
LLAMA("Llama", listOf("<EOT>")) {
CODE_LLAMA("Code Llama", listOf("<EOT>")) {
override fun buildPrompt(prefix: String, suffix: String): String {
return "<PRE> $prefix <SUF>$suffix <MID>"
}
},
CODE_GEMMA(
"CodeGemma Instruct",
listOf("<|file_separator|>", "<|fim_prefix|>", "<|fim_suffix|>", "<|fim_middle|>", "<eos>")
) {
override fun buildPrompt(prefix: String, suffix: String): String {
return "<|fim_prefix|>$prefix<|fim_suffix|>$suffix<|fim_middle|>"
}
},
STABILITY("Stability AI", listOf("<|endoftext|>")) {
override fun buildPrompt(prefix: String, suffix: String): String {
return "<fim_prefix>$prefix<fim_suffix>$suffix<fim_middle>"

View file

@ -37,7 +37,7 @@ class CustomServiceCodeCompletionForm(state: CustomServiceCodeCompletionSettings
private val promptTemplateComboBox =
ComboBox(EnumComboBoxModel(InfillPromptTemplate::class.java)).apply {
selectedItem = state.infillTemplate
setSelectedItem(InfillPromptTemplate.LLAMA)
setSelectedItem(InfillPromptTemplate.CODE_LLAMA)
addItemListener {
updatePromptTemplateHelpTooltip(it.item as InfillPromptTemplate)
}