Skip to content

Commit

Permalink
Merge pull request #206 from rjmacarthy/development
Browse files Browse the repository at this point in the history
development > main
  • Loading branch information
rjmacarthy authored Apr 9, 2024
2 parents f3560c7 + f084e47 commit 0392530
Show file tree
Hide file tree
Showing 5 changed files with 30 additions and 15 deletions.
4 changes: 2 additions & 2 deletions package-lock.json

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion package.json
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
12 changes: 3 additions & 9 deletions src/common/constants.ts
Original file line number Diff line number Diff line change
Expand Up @@ -91,6 +91,7 @@ export const FIM_TEMPLATE_FORMAT = {
llama: 'llama',
stableCode: 'stable-code',
starcoder: 'starcoder',
codegemma: 'codegemma',
custom: 'custom-template'
}

Expand All @@ -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 = [
Expand All @@ -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',
Expand Down
23 changes: 22 additions & 1 deletion src/extension/fim-templates.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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'
Expand Down Expand Up @@ -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) ||
Expand All @@ -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)
}

Expand All @@ -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)
}

Expand Down Expand Up @@ -143,6 +159,10 @@ export const getStopWordsAuto = (fimModel: string) => {
return ['<|endoftext|>']
}

if (fimModel.includes(FIM_TEMPLATE_FORMAT.codegemma)) {
return STOP_CODEGEMMA
}

return STOP_LLAMA
}

Expand All @@ -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
}

Expand Down
4 changes: 2 additions & 2 deletions src/webview/providers.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -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 = () => {
Expand Down Expand Up @@ -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>
Expand Down

0 comments on commit 0392530

Please sign in to comment.