add more stop tokens for qwen and starcoder

twinnydotdev · Sep 25, 2024 · 1064518 · 1064518
1 parent 0dba01f
commit 1064518
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 14 deletions.
diff --git a/src/common/constants.ts b/src/common/constants.ts
@@ -80,7 +80,7 @@ export const EVENT_NAME = {
   twinnyStopGeneration: 'twinny-stop-generation',
   twinnyStopSymmetryProvider: 'twinny-stop-symmetry-provider',
   twinnySymmetryModeles: 'twinny-symmetry-models',
-  twinnyTextSelection: 'twinny-text-selection',
+  twinnyTextSelection: 'twinny-text-selection'
 }
 
 export const TWINNY_COMMAND_NAME = {
@@ -156,7 +156,7 @@ export const WORKSPACE_STORAGE_KEY = {
   showEmbeddingOptions: 'showEmbeddingOptions',
   showProviders: 'showProviders',
   reviewOwner: 'reviewOwner',
-  reviewRepo: 'reviewRepo',
+  reviewRepo: 'reviewRepo'
 }
 
 export const EXTENSION_SETTING_KEY = {
@@ -220,7 +220,22 @@ export const STOP_DEEPSEEK = [
   '<｜end▁of▁sentence｜>'
 ]
 
-export const STOP_STARCODER_QWEN = ['<file_sep>', '<fim_prefix>', '<|file_sep|>', '<|fim_prefix|>', '<|endoftext|>']
+export const STOP_STARCODER = [
+  '<|endoftext|>',
+  '<file_sep>',
+  '<file_sep>',
+  '<fim_prefix>',
+  '<repo_name>',
+]
+
+export const STOP_QWEN = [
+  '<|endoftext|>',
+  '<|file_sep|>',
+  '<|fim_prefix|>',
+  '<|im_end|>',
+  '<|im_start|>',
+  '<|repo_name|>',
+]
 
 export const STOP_CODEGEMMA = ['<|file_separator|>', '<|end_of_turn|>', '<eos>']
 

diff --git a/src/extension/fim-templates.ts b/src/extension/fim-templates.ts
@@ -1,10 +1,11 @@
 import {
   FIM_TEMPLATE_FORMAT,
-  STOP_DEEPSEEK,
-  STOP_LLAMA,
-  STOP_STARCODER_QWEN,
   STOP_CODEGEMMA,
   STOP_CODESTRAL,
+  STOP_DEEPSEEK,
+  STOP_LLAMA,
+  STOP_QWEN,
+  STOP_STARCODER,
 } from '../common/constants'
 import { supportedLanguages } from '../common/languages'
 import { FimPromptTemplate } from '../common/types'
@@ -80,7 +81,7 @@ export const getFimPromptTemplateCodestral = ({
 }
 
 export const getFimPromptTemplateQwen = ({
-  prefixSuffix,
+  prefixSuffix
 }: FimPromptTemplate) => {
   const { prefix, suffix } = prefixSuffix
   return `<|fim_prefix|>${prefix}<|fim_suffix|>${suffix}<|fim_middle|>`
@@ -187,10 +188,13 @@ export const getStopWordsAuto = (fimModel: string) => {
 
   if (
     fimModel.includes(FIM_TEMPLATE_FORMAT.stableCode) ||
-    fimModel.includes(FIM_TEMPLATE_FORMAT.starcoder) ||
-    fimModel.includes(FIM_TEMPLATE_FORMAT.codeqwen)
+    fimModel.includes(FIM_TEMPLATE_FORMAT.starcoder)
   ) {
-    return STOP_STARCODER_QWEN
+    return STOP_STARCODER
+  }
+
+  if (fimModel.includes(FIM_TEMPLATE_FORMAT.codeqwen)) {
+    return STOP_QWEN
   }
 
   if (fimModel.includes(FIM_TEMPLATE_FORMAT.codegemma)) {
@@ -207,14 +211,14 @@ export const getStopWordsAuto = (fimModel: string) => {
 export const getStopWordsChosen = (format: string) => {
   if (format === FIM_TEMPLATE_FORMAT.codellama) return STOP_LLAMA
   if (format === FIM_TEMPLATE_FORMAT.deepseek) return STOP_DEEPSEEK
-  if (format === FIM_TEMPLATE_FORMAT.codeqwen) return STOP_STARCODER_QWEN
+  if (format === FIM_TEMPLATE_FORMAT.codeqwen) return STOP_QWEN
+  if (format === FIM_TEMPLATE_FORMAT.codegemma) return STOP_CODEGEMMA
+  if (format === FIM_TEMPLATE_FORMAT.codestral) return STOP_CODESTRAL
   if (
     format === FIM_TEMPLATE_FORMAT.stableCode ||
     format === FIM_TEMPLATE_FORMAT.starcoder
   )
-    return STOP_STARCODER_QWEN
-  if (format === FIM_TEMPLATE_FORMAT.codegemma) return STOP_CODEGEMMA
-  if (format === FIM_TEMPLATE_FORMAT.codestral) return STOP_CODESTRAL
+    return STOP_STARCODER
   return STOP_LLAMA
 }