fim fim for lmstudio

twinnydotdev · Apr 4, 2024 · 199ac65 · 199ac65
1 parent bacab06
commit 199ac65
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 5 deletions.
diff --git a/src/extension/provider-options.ts b/src/extension/provider-options.ts
@@ -56,7 +56,7 @@ export function createStreamRequestBodyFim(
   prompt: string,
   options: {
     temperature: number
-    numPredictChat: number
+    numPredictFim: number
     model: string
     keepAlive?: string | number
   }
@@ -71,23 +71,30 @@ export function createStreamRequestBodyFim(
         keep_alive: options.keepAlive,
         options: {
           temperature: options.temperature,
-          num_predict: options.numPredictChat
+          num_predict: options.numPredictFim
         }
       }
+    case ApiProviders.LMStudio:
+      return {
+        prompt,
+        stream: true,
+        temperature: options.temperature,
+        n_predict: options.numPredictFim
+      }
     case ApiProviders.LlamaCpp:
       return {
         prompt,
         stream: true,
         temperature: options.temperature,
-        n_predict: options.numPredictChat
+        n_predict: options.numPredictFim
       }
     case ApiProviders.LiteLLM:
     default:
       return {
         messages: [{ content: prompt, role: USER }],
         model: options.model,
         stream: true,
-        max_tokens: options.numPredictChat,
+        max_tokens: options.numPredictFim,
         temperature: options.temperature
       }
   }

diff --git a/src/extension/providers/completion.ts b/src/extension/providers/completion.ts
@@ -200,7 +200,7 @@ export class CompletionProvider implements InlineCompletionItemProvider {
       prompt,
       {
         model: this._fimModel,
-        numPredictChat: this._numPredictFim,
+        numPredictFim: this._numPredictFim,
         temperature: this._temperature,
         keepAlive: this._keepAlive
       }