diff --git a/src/config.ts b/src/config.ts index a6294b1..4530a88 100644 --- a/src/config.ts +++ b/src/config.ts @@ -33,6 +33,8 @@ class Config { modelFormat = 'deepseek'; } else if (modelName.startsWith('stable-code')) { modelFormat = 'stable-code'; + } else if (modelName.startsWith('codegemma')) { + modelFormat = 'codegemma'; } } diff --git a/src/prompts/processors/models.ts b/src/prompts/processors/models.ts index 058905f..5e12837 100644 --- a/src/prompts/processors/models.ts +++ b/src/prompts/processors/models.ts @@ -1,4 +1,4 @@ -export type ModelFormat = 'codellama' | 'deepseek' | 'stable-code'; +export type ModelFormat = 'codellama' | 'deepseek' | 'stable-code' | 'codegemma'; export function adaptPrompt(args: { format: ModelFormat, prefix: string, suffix: string }): { prompt: string, stop: string[] } { @@ -26,6 +26,13 @@ export function adaptPrompt(args: { format: ModelFormat, prefix: string, suffix: }; } + // Codegemma FIM + if (args.format === 'codegemma') { + return { + prompt: `<|fim_prefix|>${args.prefix}<|fim_suffix|>${args.suffix}<|fim_middle|>`, + stop: ['<|file_separator|>'] + } + } // Codellama FIM return { prompt: `
 ${args.prefix}  ${args.suffix} `,