From 5b3b1fa7b9278780f153b54c7a8ea73fa466a9e2 Mon Sep 17 00:00:00 2001
From: Bohan Cheng <47214785+cbh778899@users.noreply.github.com>
Date: Tue, 15 Oct 2024 10:42:05 +1100
Subject: [PATCH] fix wllama load bugs (#41)

* remove output

Signed-off-by: cbh778899 <cbh778899@outlook.com>

* rename worker.js to wllama-worker.js

Signed-off-by: cbh778899 <cbh778899@outlook.com>

* remove force load platform to fix bug

Signed-off-by: cbh778899 <cbh778899@outlook.com>

* add download protector call to wllama settings

Signed-off-by: cbh778899 <cbh778899@outlook.com>

* change import path

Signed-off-by: cbh778899 <cbh778899@outlook.com>

---------

Signed-off-by: cbh778899 <cbh778899@outlook.com>
---
 src/components/chat/index.jsx                 |  7 +----
 src/components/settings/LlamaSettings.jsx     |  1 -
 src/components/settings/WllamaSettings.jsx    | 29 ++++++++++++++++---
 src/utils/start_loader.js                     |  2 +-
 src/utils/workers/index.js                    | 10 ++-----
 .../workers/{worker.js => wllama-worker.js}   | 16 ++++++++--
 6 files changed, 44 insertions(+), 21 deletions(-)
 rename src/utils/workers/{worker.js => wllama-worker.js} (93%)

diff --git a/src/components/chat/index.jsx b/src/components/chat/index.jsx
index 65654ae..b9a3424 100644
--- a/src/components/chat/index.jsx
+++ b/src/components/chat/index.jsx
@@ -21,7 +21,7 @@ export default function Chat() {
 
     const idb = useIDB();
     // const settings = useRef(getCompletionFunctions());
-    const settings = useRef(getCompletionFunctions('Llama'));
+    const settings = useRef(getCompletionFunctions());
 
     async function sendMessage(message, files) {
         // save user messages
@@ -143,11 +143,6 @@ export default function Chat() {
                 setHistory={setTickets} history={tickets} 
                 deleteHistory={requestDelete} platform={settings.current.platform}
             />
-            {/* <Conversation 
-                uid={chat.uid} 
-                title={chat.title} updateTitle={updateTitle}
-                client={chat.client} updateClient={updateChatClient}
-            /> */}
             <ChatPage 
                 updateTitle={updateTitle}
                 chat={chat} chat_history={chat_history}
diff --git a/src/components/settings/LlamaSettings.jsx b/src/components/settings/LlamaSettings.jsx
index 7788b73..d6d163c 100644
--- a/src/components/settings/LlamaSettings.jsx
+++ b/src/components/settings/LlamaSettings.jsx
@@ -22,7 +22,6 @@ export default function LlamaSettings({ trigger, enabled, updateEnabled, openDow
         if(enabled) {
             // check if model with this url already downloaded
             let stored_model = await idb.getOne("downloaded-models", {where: [{ platform: 'Llama', url }]})
-            console.log(stored_model)
             // if no model record, means not downloaded
             if(!stored_model) {
                 await openDownloadProtector(
diff --git a/src/components/settings/WllamaSettings.jsx b/src/components/settings/WllamaSettings.jsx
index c606e1a..83de0b9 100644
--- a/src/components/settings/WllamaSettings.jsx
+++ b/src/components/settings/WllamaSettings.jsx
@@ -3,16 +3,19 @@ import SettingSection from "./SettingSection";
 import TrueFalseComponent from "./components/TrueFalseComponent";
 import ScrollBarComponent from "./components/ScrollBarComponent";
 import { getPlatformSettings, updatePlatformSettings } from "../../utils/general_settings";
-import { loadModel, loadModelSamplingSettings } from "../../utils/workers/worker";
+import { downloadModel, isModelDownloaded, loadModel, loadModelSamplingSettings } from "../../utils/workers/wllama-worker";
+import useIDB from "../../utils/idb";
 
-export default function WllamaSettings({ trigger, enabled, updateEnabled }) {
+export default function WllamaSettings({ trigger, enabled, updateEnabled, openDownloadProtector }) {
 
     const [ threads, setThreads ] = useState(1);
     const [ batch_size, setBatchSize ] = useState(256);
     const [ context_length, setContextLength ] = useState(4096);
     const [ continue_conv, setContinueConversation ] = useState(false);
 
-    function saveSettings() {
+    const idb = useIDB();
+
+    async function saveSettings() {
         updatePlatformSettings({
             wllama_threads: threads,
             wllama_batch_size: batch_size,
@@ -22,7 +25,25 @@ export default function WllamaSettings({ trigger, enabled, updateEnabled }) {
 
         if(enabled) {
             loadModelSamplingSettings();
-            loadModel('completion');
+            const model = 'completion'
+            if(!(await isModelDownloaded(model))) {
+                await openDownloadProtector(
+                    "Please wait while we downloading the model...",
+                    `Downloading model smollm-360m-instruct-add-basics-q8_0.gguf`,
+                    async callback => {
+                        const download_result = await downloadModel(model, callback);
+                        if(!download_result) return;
+                        const { model_name, url, finish_time, size } = download_result;
+                        await idb.insert("downloaded-models", {
+                            'model-name': model_name,
+                            url, size, createdAt: finish_time,
+                            platform: 'Wllama'
+                        })
+                    }
+                )
+            }
+
+            await loadModel('completion');
         }
     }
 
diff --git a/src/utils/start_loader.js b/src/utils/start_loader.js
index 974fb7a..6bc3f1a 100644
--- a/src/utils/start_loader.js
+++ b/src/utils/start_loader.js
@@ -1,6 +1,6 @@
 import { getPlatformSettings } from "./general_settings";
 import { instance } from "./idb";
-import { loadModel } from "./workers/worker";
+import { loadModel } from "./workers/wllama-worker";
 
 export default async function loader() {
     localStorage.setItem('not-first-time', '1');
diff --git a/src/utils/workers/index.js b/src/utils/workers/index.js
index 5cdeddb..a994429 100644
--- a/src/utils/workers/index.js
+++ b/src/utils/workers/index.js
@@ -1,5 +1,5 @@
 import { getModelSettings, getPlatformSettings } from "../general_settings";
-import { chatCompletions as WllamaCompletions, abortCompletion as WllamaAbort } from "./worker";
+import { chatCompletions as WllamaCompletions, abortCompletion as WllamaAbort, setClient as WllamaSetClient } from "./wllama-worker";
 import { chatCompletions as AwsCompletions, abortCompletion as AwsAbort, setClient as AwsSetClient, formator as AwsFormator } from "./aws-worker"
 import { chatCompletions as OpenaiCompletions, abortCompletion as OpenaiAbort, setClient as OpenAISetClient } from "./openai-worker";
 
@@ -45,13 +45,9 @@ export function getCompletionFunctions(platform = null) {
             return {
                 completions: WllamaCompletions,
                 abort: WllamaAbort,
-                platform: "Wllama"
+                platform: "Wllama",
+                initClient: WllamaSetClient
             }
-        // default:
-        //     return { 
-        //         completions: WllamaCompletions, abort: WllamaAbort, 
-        //         platform: "Wllama"
-        //     }
     }
     
 }
\ No newline at end of file
diff --git a/src/utils/workers/worker.js b/src/utils/workers/wllama-worker.js
similarity index 93%
rename from src/utils/workers/worker.js
rename to src/utils/workers/wllama-worker.js
index c671616..427c8e5 100644
--- a/src/utils/workers/worker.js
+++ b/src/utils/workers/wllama-worker.js
@@ -44,6 +44,7 @@ loadModelSamplingSettings();
 
 const DEFAULT_CHAT_TEMPLATE = "{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}";
 
+
 const engines = {
     completion: {
         model_src: "https://huggingface.co/HuggingFaceTB/smollm-360M-instruct-v0.2-Q8_0-GGUF/resolve/main/smollm-360m-instruct-add-basics-q8_0.gguf",
@@ -52,6 +53,11 @@ const engines = {
     }
 }
 
+export function setClient() {
+    // we don't need to reset client for wllama
+    return 'fake-client'
+}
+
 let stop_signal = false;
 
 export async function ReloadEngine(type = 'completion') {
@@ -79,10 +85,16 @@ export async function downloadModel(type = 'completion', cb = null) {
         allowOffline: true,
         embeddings: type === 'embedding',
         progressCallback: ({loaded, total})=>{
-            cb && cb((loaded / total) * 100);
+            cb && cb((loaded / total) * 100, false);
         }
     })
-    cb && cb(100);
+    cb && cb(100, true);
+    return {
+        model_name: model_src.split('/').pop(),
+        url: model_src,
+        size: 0,
+        finish_time: Date.now()
+    }
 }
 
 export async function deleteModel(type = 'completion') {