Merge pull request #206 from rjmacarthy/development

development > main

package.json

@@ -2,7 +2,7 @@
   "name": "twinny",
   "displayName": "twinny - AI Code Completion and Chat",
   "description": "Locally hosted AI code completion plugin for vscode",
-  "version": "3.11.0",
+  "version": "3.11.1",
   "icon": "assets/icon.png",
   "keywords": [
     "code-inference",

src/common/constants.ts

@@ -91,6 +91,7 @@ export const FIM_TEMPLATE_FORMAT = {
   llama: 'llama',
   stableCode: 'stable-code',
   starcoder: 'starcoder',
+  codegemma: 'codegemma',
   custom: 'custom-template'
 }
 
@@ -106,6 +107,8 @@ export const STOP_DEEPSEEK = [
 
 export const STOP_STARCODER = ['<|endoftext|>']
 
+export const STOP_CODEGEMMA = ['<|file_separator|>', '<|end_of_turn|>', '<eos>']
+
 export const DEFAULT_TEMPLATE_NAMES = defaultTemplates.map(({ name }) => name)
 
 export const DEFAULT_ACTION_TEMPLATES = [
@@ -129,15 +132,6 @@ export const DEFAULT_PROVIDER_FORM_VALUES = {
   type: 'chat'
 }
 
-export const FIM_TEMPLATE_TYPE = [
-  'automatic',
-  'stable-code',
-  'codellama',
-  'deepseek',
-  'starcoder',
-  'custom-template'
-]
-
 export const WASM_LANGAUAGES: { [key: string]: string } = {
   cpp: 'cpp',
   hpp: 'cpp',

src/extension/fim-templates.ts

@@ -2,7 +2,8 @@ import {
   FIM_TEMPLATE_FORMAT,
   STOP_DEEPSEEK,
   STOP_LLAMA,
-  STOP_STARCODER
+  STOP_STARCODER,
+  STOP_CODEGEMMA
 } from '../common/constants'
 import { supportedLanguages } from '../common/languages'
 import { FimPromptTemplate } from '../common/types'
@@ -72,6 +73,13 @@ export const getFimPromptTemplateStableCode = ({
   return `<fim_prefix>${fileContext}\n${heading}${prefix}<fim_suffix>${suffix}<fim_middle>`
 }
 
+export const getFimPromptTemplateCodegemma = ({
+  prefixSuffix
+}: FimPromptTemplate) => {
+  const { prefix, suffix } = prefixSuffix
+  return `<|fim_prefix|>${prefix}<|fim_suffix|>${suffix}<|fim_middle|>`
+}
+
 function getFimTemplateAuto(fimModel: string, args: FimPromptTemplate) {
   if (
     fimModel.includes(FIM_TEMPLATE_FORMAT.codellama) ||
@@ -91,6 +99,10 @@ function getFimTemplateAuto(fimModel: string, args: FimPromptTemplate) {
     return getFimPromptTemplateStableCode(args)
   }
 
+  if (fimModel.includes(FIM_TEMPLATE_FORMAT.codegemma)) {
+    return getFimPromptTemplateCodegemma(args)
+  }
+
   return getDefaultFimPromptTemplate(args)
 }
 
@@ -110,6 +122,10 @@ function getFimTemplateChosen(format: string, args: FimPromptTemplate) {
     return getFimPromptTemplateStableCode(args)
   }
 
+  if (format === FIM_TEMPLATE_FORMAT.codegemma) {
+    return getFimPromptTemplateCodegemma(args)
+  }
+
   return getDefaultFimPromptTemplate(args)
 }
 
@@ -143,6 +159,10 @@ export const getStopWordsAuto = (fimModel: string) => {
     return ['<|endoftext|>']
   }
 
+  if (fimModel.includes(FIM_TEMPLATE_FORMAT.codegemma)) {
+    return STOP_CODEGEMMA
+  }
+
   return STOP_LLAMA
 }
 
@@ -154,6 +174,7 @@ export const getStopWordsChosen = (format: string) => {
     format === FIM_TEMPLATE_FORMAT.starcoder
   )
     return STOP_STARCODER
+  if (format === FIM_TEMPLATE_FORMAT.codegemma) return STOP_CODEGEMMA
   return STOP_LLAMA
 }
 

src/webview/providers.tsx

@@ -14,7 +14,7 @@ import styles from './providers.module.css'
 import { TwinnyProvider } from '../extension/provider-manager'
 import {
   DEFAULT_PROVIDER_FORM_VALUES,
-  FIM_TEMPLATE_TYPE
+  FIM_TEMPLATE_FORMAT,
 } from '../common/constants'
 
 export const Providers = () => {
@@ -223,7 +223,7 @@ function ProviderForm({ onClose, provider }: ProviderFormProps) {
               onChange={handleChangeDropdown}
               value={formState.fimTemplate}
             >
-              {FIM_TEMPLATE_TYPE.map((type, index) => (
+              {Object.values(FIM_TEMPLATE_FORMAT).map((type, index) => (
                 <VSCodeOption key={index} value={type || 'automatic'}>
                   {type}
                 </VSCodeOption>