freeze demo using model revisions

src/hooks/useLLM.ts

@@ -5,6 +5,21 @@ import {
   TextStreamer,
 } from "@huggingface/transformers";
 
+// Define the supported model IDs
+export type SupportedModelId = "350M" | "700M" | "1.2B";
+
+export interface ModelConfig {
+  dtype: string;
+  device: string;
+  revision?: string;
+}
+
+export const MODEL_CONFIGS: Record<SupportedModelId, ModelConfig> = {
+  "350M": { dtype: "q4f16", device: "webgpu", revision: "5bc4b3e8cfd21660c0b1b9faa447ffbd9926b829" },
+  "700M": { dtype: "q4f16", device: "webgpu", revision: "bf72eeabfe73a798674db899830a0dca99f8eabc" },
+  "1.2B": { dtype: "q4f16", device: "webgpu", revision: "7f871660813dc1f34f0d304c77506c5fbdb440a0" },
+};
+
 interface LLMState {
   isLoading: boolean;
   isReady: boolean;
@@ -24,7 +39,7 @@ let moduleCache: {
   };
 } = {};
 
-export const useLLM = (modelId?: string) => {
+export const useLLM = (modelId?: SupportedModelId | string) => {
   const [state, setState] = useState<LLMState>({
     isLoading: false,
     isReady: false,
@@ -105,15 +120,37 @@ export const useLLM = (modelId?: string) => {
           }
         };
 
-        const tokenizer = await AutoTokenizer.from_pretrained(MODEL_ID, {
-          progress_callback: progressCallback,
-        });
-
-        const model = await AutoModelForCausalLM.from_pretrained(MODEL_ID, {
+        // Fallback to defaults if an unknown modelId string is passed
+        const config = MODEL_CONFIGS[modelId as SupportedModelId] || {
           dtype: "q4f16",
           device: "webgpu",
+        };
+
+        const tokenizerOptions: Record<string, any> = {
+          progress_callback: progressCallback,
+        };
+        if (config.revision) {
+          tokenizerOptions.revision = config.revision;
+        }
+
+        const tokenizer = await AutoTokenizer.from_pretrained(
+          MODEL_ID,
+          tokenizerOptions
+        );
+
+        const modelOptions: Record<string, any> = {
+          dtype: config.dtype,
+          device: config.device,
           progress_callback: progressCallback,
-        });
+        };
+        if (config.revision) {
+          modelOptions.revision = config.revision;
+        }
+
+        const model = await AutoModelForCausalLM.from_pretrained(
+          MODEL_ID,
+          modelOptions
+        );
 
         const instance = { model, tokenizer };
         instanceRef.current = instance;
@@ -231,4 +268,4 @@ export const useLLM = (modelId?: string) => {
     clearPastKeyValues,
     cleanup,
   };
-};
+};