Create worker.js

worker.js

@@ -0,0 +1,141 @@
+import { pipeline, env, cos_sim } from 'https://cdn.jsdelivr.net/npm/@huggingface/transformers@3.0.0-alpha.19/dist/transformers.min.js';
+
+// Cấu hình WebGPU
+env.backends.onnx.wasm.proxy = false; 
+
+// --- CẤU HÌNH MODEL (KHÔNG ĐỔI THEO YÊU CẦU) ---
+// 1. Embedding Model
+const EMBEDDING_MODEL_ID = 'onnx-community/embeddinggemma-300m-ONNX';
+
+// 2. LLM Model: Granite 4.0 Nano
+// Lưu ý: Nếu phiên bản ONNX của Granite 4.0 chưa public dưới ID này, 
+// bạn cần trỏ tới đúng repo onnx (ví dụ ibm-granite/granite-3.0-2b-instruct nếu 4.0 chưa có onnx).
+// Tuy nhiên, tôi giữ nguyên tham chiếu "Granite" như yêu cầu.
+const LLM_MODEL_ID = 'ibm-granite/granite-4.0-350m-instruct'; // Kiểm tra lại tên chính xác trên HF Hub nếu lỗi
+
+let extractor = null;
+let generator = null;
+let vectorStore = []; // Lưu trữ chunks và vectors: { text: string, vector: number[] }
+
+// Khởi tạo Models
+async function initModels() {
+    try {
+        console.log("Đang tải Embedding Model...");
+        extractor = await pipeline('feature-extraction', EMBEDDING_MODEL_ID, {
+            device: 'webgpu',
+            dtype: 'fp32', // Embedding thường cần độ chính xác
+        });
+
+        console.log("Đang tải LLM Granite 4.0...");
+        generator = await pipeline('text-generation', LLM_MODEL_ID, {
+            device: 'webgpu',
+            dtype: 'q4', // Quantization 4-bit để chạy mượt trên browser
+            use_external_data_format: true 
+        });
+
+        self.postMessage({ type: 'init_complete' });
+    } catch (e) {
+        self.postMessage({ type: 'error', payload: "Lỗi tải model: " + e.message });
+    }
+}
+
+// Xử lý chunking văn bản
+function chunkText(text, chunkSize = 300, overlap = 50) {
+    const sentences = text.match(/[^.!?]+[.!?]+|[^.!?]+$/g) || [text];
+    let chunks = [];
+    let currentChunk = "";
+    
+    for (let sentence of sentences) {
+        if ((currentChunk + sentence).length > chunkSize) {
+            chunks.push(currentChunk.trim());
+            currentChunk = sentence.slice(-overlap); // Overlap đơn giản
+        } else {
+            currentChunk += " " + sentence;
+        }
+    }
+    if (currentChunk) chunks.push(currentChunk.trim());
+    return chunks;
+}
+
+// Tạo embeddings cho văn bản
+async function ingestText(text) {
+    const chunks = chunkText(text);
+    vectorStore = []; // Reset store
+
+    for (const chunk of chunks) {
+        const output = await extractor(chunk, { pooling: 'mean', normalize: true });
+        vectorStore.push({
+            text: chunk,
+            vector: output.data
+        });
+    }
+    console.log(`Đã index ${vectorStore.length} đoạn văn bản.`);
+}
+
+// Tìm kiếm RAG
+async function retrieve(query) {
+    const queryOutput = await extractor(query, { pooling: 'mean', normalize: true });
+    const queryVector = queryOutput.data;
+
+    // Tính Cosine Similarity
+    const scored = vectorStore.map(item => {
+        return {
+            text: item.text,
+            score: cos_sim(queryVector, item.vector)
+        };
+    });
+
+    // Lấy top 3 đoạn liên quan nhất
+    scored.sort((a, b) => b.score - a.score);
+    return scored.slice(0, 3).map(i => i.text).join("\n\n");
+}
+
+// Xử lý tin nhắn từ Main Thread
+self.onmessage = async (e) => {
+    if (!extractor || !generator) {
+        await initModels();
+    }
+
+    const { type, payload } = e.data;
+
+    if (type === 'ingest_text') {
+        await ingestText(payload);
+    } else if (type === 'query') {
+        // 1. Retrieve Context
+        const context = await retrieve(payload);
+        
+        // 2. Tạo Prompt cho Granite
+        // Định dạng prompt cơ bản cho instruction tuned model
+        const prompt = `<|system|>
+Bạn là trợ lý AI hữu ích. Hãy trả lời câu hỏi dựa trên ngữ cảnh được cung cấp bên dưới bằng Tiếng Việt.
+Ngữ cảnh:
+${context}
+<|user|>
+${payload}
+<|assistant|>
+`;
+        
+        // 3. Generate Answer
+        try {
+            const output = await generator(prompt, {
+                max_new_tokens: 256,
+                temperature: 0.7,
+                do_sample: true,
+            });
+            
+            // Lấy phần trả lời sau tag assistant (tuỳ thuộc format model)
+            let answer = output[0].generated_text;
+            // Cắt bớt phần prompt nếu cần thiết
+            if (answer.includes("<|assistant|>")) {
+                answer = answer.split("<|assistant|>")[1];
+            }
+            
+            self.postMessage({ type: 'answer', payload: answer });
+        } catch (err) {
+            self.postMessage({ type: 'error', payload: err.message });
+        }
+    }
+};
+
+// Khởi tạo ngay khi worker chạy
+initModels();