hugging2021/local-gemma4-rag

1import { pipeline, env } from 'https://cdn.jsdelivr.net/npm/@huggingface/transformers@4.0.1?v=401';2 3// Configure environment4env.allowLocalModels = false;5env.useBrowserCache = true;6 7// Define absolute path for local wasm hosting8const baseUrl = self.location.origin + self.location.pathname.substring(0, self.location.pathname.lastIndexOf('/') + 1);9env.backends.onnx.wasm.wasmPaths = baseUrl + 'wasm/';10env.backends.onnx.wasm.proxy = false;11 12let generatorPipeline = null;13let embeddingPipeline = null;14 15const EMBEDDING_MODEL = 'Xenova/all-MiniLM-L6-v2';16let CURRENT_LLM_MODEL = 'onnx-community/gemma-4-E2B-it-ONNX'; // Default17 18self.onmessage = async (e) => {19  const { action, payload } = e.data;20 21  if (action === 'init') {22    if (payload && payload.modelId) {23        CURRENT_LLM_MODEL = payload.modelId;24    }25    await initModels();26  } else if (action === 'embed') {27    const vector = await embedText(payload.text);28    self.postMessage({ action: 'embed_result', payload: { vector } });29  } else if (action === 'generate') {30    await generateResponse(payload.prompt, payload.context);31  } else if (action === 'cleanup') {32    await cleanup();33  }34};35 36async function cleanup() {37  self.postMessage({ action: 'status', payload: { text: 'Disposing AI models...' } });38  if (generatorPipeline) {39    await generatorPipeline.dispose();40    generatorPipeline = null;41  }42  if (embeddingPipeline) {43    await embeddingPipeline.dispose();44    embeddingPipeline = null;45  }46  self.postMessage({ action: 'status', payload: { text: 'Memory cleared.' } });47}48 49async function initModels() {50  const modelName = CURRENT_LLM_MODEL.split('/').pop().replace('-ONNX', '');51  self.postMessage({ action: 'status', payload: { text: `Loading Embedding Model...`, progress: 0 } });52 53  // 1. Load Embedding Model54  embeddingPipeline = await pipeline('feature-extraction', EMBEDDING_MODEL, {55    device: 'webgpu',56    progress_callback: (p) => {57      if (p.status === 'progress') {58        self.postMessage({ action: 'progress', payload: { model: 'embedding', progress: p.progress } });59      }60    }61  });62 63  self.postMessage({ action: 'status', payload: { text: `Loading ${modelName}...`, progress: 50 } });64 65  // 2. Load Selected Gemma LLM66  generatorPipeline = await pipeline('text-generation', CURRENT_LLM_MODEL, {67    device: 'webgpu',68    dtype: 'q4',69    progress_callback: (p) => {70      if (p.status === 'progress') {71        self.postMessage({ action: 'progress', payload: { model: 'llm', progress: p.progress } });72      }73    }74  });75 76  // 3. GPU Warmup (compile shaders)77  self.postMessage({ action: 'status', payload: { text: 'Warming up GPU...', progress: 95 } });78  await generatorPipeline('warmup', { max_new_tokens: 1 });79 80  self.postMessage({ action: 'ready' });81}82 83async function embedText(text) {84  const output = await embeddingPipeline(text, { pooling: 'mean', normalize: true });85  return Array.from(output.data);86}87 88async function generateResponse(userPrompt, context) {89  // Use official chat template90  const isGeneralKnowledge = !context || context === 'No relevant context found in local documents.';91  92  const systemInstruction = isGeneralKnowledge 93    ? "Act as a helpful AI assistant. Answer using your own knowledge."94    : `Answer based on the following context. If the answer is not in the context, you MAY use your own knowledge but clearly state that the information was not found in the documents.\n\nCONTEXT:\n${context}`;95 96  const messages = [97    { 98      role: 'user', 99      content: `${systemInstruction}\n\nQUESTION:\n${userPrompt}` 100    }101  ];102 103  const fullPrompt = generatorPipeline.tokenizer.apply_chat_template(messages, {104    tokenize: false,105    add_generation_prompt: true,106  });107 108  const output = await generatorPipeline(fullPrompt, {109    max_new_tokens: 1024,110    do_sample: false,111    repetition_penalty: 1.2,112    return_full_text: false,113    stop_sequences: ["<turn|>", "<channel|>", "<eos>", "<|turn|>"],114    callback_function: (beams) => {115      const decoded = generatorPipeline.tokenizer.decode(beams[0].output_token_ids, {116        skip_special_tokens: true,117      });118      119      let textToPush = decoded;120      121      // Robust Reasoning Filter (Gemma 4 specific)122      // The model often starts with "thought\n..."123      if (textToPush.toLowerCase().includes('thought')) {124        const thoughtParts = textToPush.split(/\n\n|Answer:/i);125        if (thoughtParts.length > 1) {126            textToPush = thoughtParts.slice(1).join('\n\n').trim();127        } else {128            textToPush = '_Thinking..._';129        }130      }131 132      // Cleanup any leaked turn markers (should be handled by stop_sequences but just in case)133      textToPush = textToPush.replace(/<\|turn\|>model\n/g, '').replace(/<turn\|>/g, '');134      135      self.postMessage({ action: 'chunk', payload: { text: textToPush } });136    }137  });138 139  const finalResult = output[0].generated_text;140  self.postMessage({ action: 'generate_complete', payload: { text: finalResult } });141}142