add delete all in cache

mlc-ai · Feb 26, 2024 · dcd26dc · dcd26dc
1 parent 3b0cb09
commit dcd26dc
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 10 deletions.
diff --git a/src/chat_module.ts b/src/chat_module.ts
@@ -6,7 +6,8 @@ import {
   AppConfig,
   prebuiltAppConfig,
   GenerationConfig,
-  postInitAndCheckGenerationConfigValues
+  postInitAndCheckGenerationConfigValues,
+  ModelRecord
 } from "./config";
 import { LLMChatPipeline } from "./llm_chat"
 import {
@@ -660,18 +661,46 @@ export class ChatRestModule implements ChatInterface {
   }
 }
 
+export async function findModelRecord (localId: string, appConfig?: AppConfig){
+  const matchedItem = appConfig?.model_list.find(
+    item => item.local_id == localId
+  );
+  if (matchedItem !== undefined) {
+    return matchedItem;
+  }
+  throw Error("Cannot find model_url for " + localId);
+}
+
 export async function hasModelInCache(localId: string, appConfig?: AppConfig): Promise<boolean> {
   if (appConfig === undefined) {
     appConfig = prebuiltAppConfig;
   }
-  const findModelRecord = () => {
-    const matchedItem = appConfig?.model_list.find(
-      item => item.local_id == localId
-    );
-    if (matchedItem !== undefined) return matchedItem;
-    throw Error("Cannot find model_url for " + localId);
-  }
-  const modelRecord = findModelRecord();
+  const modelRecord = await findModelRecord(localId, appConfig);
   const modelUrl = modelRecord.model_url;
   return tvmjs.hasNDArrayInCache(modelUrl, "webllm/model");
 }
+
+export async function deleteAllModelDataInCache(localId: string, appConfig?: AppConfig) {
+  if (appConfig === undefined) {
+    appConfig = prebuiltAppConfig;
+  }
+  // delete model
+  const modelRecord = await findModelRecord(localId, appConfig);
+  const modelUrl = modelRecord.model_url;
+  tvmjs.deleteNDArrayCache(modelUrl, "webllm/model");
+  const modelCache = new tvmjs.ArtifactCache("webllm/model");
+  // delete tokenizer
+  const url = new URL("tokenizer.model", modelUrl).href;
+  await modelCache.deleteInCache(url);
+  const url2 = new URL("tokenizer.json", modelUrl).href;
+  await modelCache.deleteInCache(url2);
+  // delete wasm
+  const wasmUrl = modelRecord.model_lib_url;
+  const wasmCache = new tvmjs.ArtifactCache("webllm/wasm");
+  await wasmCache.deleteInCache(wasmUrl);
+  // delete config 
+  const configCache = new tvmjs.ArtifactCache("webllm/config");
+  const configUrl = new URL("mlc-chat-config.json", modelUrl).href;
+  await configCache.deleteInCache(configUrl);
+}
+
diff --git a/src/index.ts b/src/index.ts
@@ -12,7 +12,7 @@ export {
 
 export {
   ChatModule,
-  ChatRestModule, hasModelInCache
+  ChatRestModule, hasModelInCache, deleteModelInCache
 } from "./chat_module";
 
 export {