From 38a77bbbb01ec6216c1cde34d6f9a38d245ba3b8 Mon Sep 17 00:00:00 2001
From: Steve Korshakov <steve@korshakov.com>
Date: Tue, 16 Jan 2024 15:52:03 -0800
Subject: [PATCH] feat: stable-code support, cancel download only for specific
 model

---
 package.json                     |  4 +++-
 src/config.ts                    |  7 +++++--
 src/prompts/processors/models.ts | 10 +++++++++-
 src/prompts/provider.ts          |  4 ++--
 4 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/package.json b/package.json
index 25d18f7..a6678e1 100644
--- a/package.json
+++ b/package.json
@@ -45,6 +45,7 @@
           "inference.model": {
             "type": "string",
             "enum": [
+              "stable-code:3b-code-q4_0",
               "codellama:7b-code-q4_K_S",
               "codellama:7b-code-q4_K_M",
               "codellama:7b-code-q6_K",
@@ -70,7 +71,7 @@
               "deepseek-coder:33b-base-fp16",
               "custom"
             ],
-            "default": "deepseek-coder:1.3b-base-q4_1",
+            "default": "stable-code:3b-code-q4_0",
             "description": "Inference model to use",
             "order": 2
           },
@@ -89,6 +90,7 @@
           "inference.custom.format": {
             "type": "string",
             "enum": [
+              "stable-code",
               "codellama",
               "deepseek"
             ],
diff --git a/src/config.ts b/src/config.ts
index 30fefbc..bb10809 100644
--- a/src/config.ts
+++ b/src/config.ts
@@ -1,4 +1,5 @@
 import vscode from 'vscode';
+import { ModelFormat } from './prompts/processors/models';
 
 class Config {
 
@@ -22,13 +23,15 @@ class Config {
 
         // Load model
         let modelName = config.get('model') as string;
-        let modelFormat: 'codellama' | 'deepseek' = 'codellama';
+        let modelFormat: ModelFormat = 'codellama';
         if (modelName === 'custom') {
             modelName = config.get('custom.model') as string;
-            modelFormat = config.get('cutom.format') as 'codellama' | 'deepseek';
+            modelFormat = config.get('cutom.format') as ModelFormat;
         } else {
             if (modelName.startsWith('deepseek-coder')) {
                 modelFormat = 'deepseek';
+            } else if (modelName.startsWith('stable-code')) {
+                modelFormat = 'stable-code';
             }
         }
 
diff --git a/src/prompts/processors/models.ts b/src/prompts/processors/models.ts
index 246bb55..6133028 100644
--- a/src/prompts/processors/models.ts
+++ b/src/prompts/processors/models.ts
@@ -1,4 +1,4 @@
-export type ModelFormat = 'codellama' | 'deepseek';
+export type ModelFormat = 'codellama' | 'deepseek' | 'stable-code';
 
 export function adaptPrompt(args: { format: ModelFormat, prefix: string, suffix: string }): { prompt: string, stop: string[] } {
 
@@ -18,6 +18,14 @@ export function adaptPrompt(args: { format: ModelFormat, prefix: string, suffix:
         };
     }
 
+    // Stable code FIM
+    if (args.format === 'stable-code')  {
+        return {
+            prompt: `<fim_prefix>${args.prefix}<fim_suffix>${args.suffix}<fim_middle>`,
+            stop: [`<|endoftext|>`]
+        };
+    }
+
     // Codellama FIM
     return {
         prompt: `<PRE> ${args.prefix} <SUF>${args.suffix} <MID>`,
diff --git a/src/prompts/provider.ts b/src/prompts/provider.ts
index 5f37360..0e03c95 100644
--- a/src/prompts/provider.ts
+++ b/src/prompts/provider.ts
@@ -91,7 +91,7 @@ export class PromptProvider implements vscode.InlineCompletionItemProvider {
                         if (!modelExists) {
 
                             // Check if user asked to ignore download
-                            if (this.context.globalState.get('llama-coder-download-ignored')) {
+                            if (this.context.globalState.get('llama-coder-download-ignored') === inferenceConfig.modelName) {
                                 info(`Ingoring since user asked to ignore download.`);
                                 return;
                             }
@@ -100,7 +100,7 @@ export class PromptProvider implements vscode.InlineCompletionItemProvider {
                             let download = await vscode.window.showInformationMessage(`Model ${inferenceConfig.modelName} is not downloaded. Do you want to download it? Answering "No" would require you to manually download model.`, 'Yes', 'No');
                             if (download === 'No') {
                                 info(`Ingoring since user asked to ignore download.`);
-                                this.context.globalState.update('llama-coder-download-ignored', true);
+                                this.context.globalState.update('llama-coder-download-ignored', inferenceConfig.modelName);
                                 return;
                             }