Update app.py

app.py

@@ -1,58 +1,119 @@
+# --- RAG / Semantic Search imports ---
 import numpy as np
 import traceback
 import torch
+import time
+
 from langchain_text_splitters import MarkdownHeaderTextSplitter, RecursiveCharacterTextSplitter
 from langchain_huggingface import HuggingFaceEmbeddings
-from openai import OpenAI
-import time
-import os 
 
-api_key = os.environ.get("GEMINI_API_KEY")
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-if not api_key:
-    raise ValueError("❌ Lỗi: Không tìm thấy GEMINI_API_KEY. Vui lòng cấu hình trong Settings -> Secrets.")
+# =========================================================
+# LLM CONFIG (HuggingFace local)
+# =========================================================
+LLM_NAME = "Qwen/Qwen2.5-7B-Instruct"   # đổi nếu GPU mạnh hơn
+DTYPE = torch.float16 if torch.cuda.is_available() else torch.float32
 
-client = OpenAI(
-    api_key=api_key,
-    base_url="https://generativelanguage.googleapis.com/v1beta/openai/"
+tokenizer = AutoTokenizer.from_pretrained(LLM_NAME)
+llm_model = AutoModelForCausalLM.from_pretrained(
+    LLM_NAME,
+    torch_dtype=DTYPE,
+    device_map="auto"
 )
+llm_model.eval()
+
+
+def _llm_generate(prompt, max_new_tokens=256):
+    inputs = tokenizer(prompt, return_tensors="pt").to(llm_model.device)
+    with torch.no_grad():
+        outputs = llm_model.generate(
+            **inputs,
+            max_new_tokens=max_new_tokens,
+            do_sample=False,
+            temperature=0
+        )
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+
+# =========================================================
+# --- Functions for RAG (GIỮ NGUYÊN TÊN HÀM)
+# =========================================================
 
 def md_to_kb_safe(md_text, embedding_model_name="sentence-transformers/all-MiniLM-L6-v2"):
     try:
-        headers_to_split_on = [("#", "Header 1"), ("##", "Header 2"), ("###", "Header 3")]
+        headers_to_split_on = [
+            ("#", "Header 1"),
+            ("##", "Header 2"),
+            ("###", "Header 3"),
+        ]
         splitter = MarkdownHeaderTextSplitter(headers_to_split_on=headers_to_split_on)
         md_chunks = splitter.split_text(md_text)
-        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200, length_function=len)
+
+        text_splitter = RecursiveCharacterTextSplitter(
+            chunk_size=1000,
+            chunk_overlap=200,
+            length_function=len,
+        )
         final_chunks = text_splitter.split_documents(md_chunks)
+
         texts = [doc.page_content for doc in final_chunks]
-        device = "cuda" if torch.cuda.is_available() and torch.cuda.memory_allocated() < 2_000_000_000 else "cpu"
-        embedding_model = HuggingFaceEmbeddings(model_name=embedding_model_name, model_kwargs={"device": device})
+
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        embedding_model = HuggingFaceEmbeddings(
+            model_name=embedding_model_name,
+            model_kwargs={"device": device},
+        )
+
         vectors = embedding_model.embed_documents(texts)
+
         kb = [{"text": texts[i], "vector": vectors[i]} for i in range(len(texts))]
-        return {"success": True, "num_chunks": len(final_chunks), "kb": kb, "embed_model": embedding_model}
+
+        return {
+            "success": True,
+            "num_chunks": len(final_chunks),
+            "kb": kb,
+            "embed_model": embedding_model,
+        }
+
     except Exception as e:
-        return {"success": False, "error": str(e), "traceback": traceback.format_exc()}
+        return {
+            "success": False,
+            "error": str(e),
+            "traceback": traceback.format_exc(),
+        }
+
 
 def cosine_similarity(v1, v2):
     return np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2))
 
+
 def semantic_search(query, embed_model, kb, top_k=3):
     t0 = time.time()
     q_vec = np.array(embed_model.embed_query(query))
-    scores = [(cosine_similarity(q_vec, item["vector"]), item["text"]) for item in kb]
+
+    scores = [
+        (cosine_similarity(q_vec, item["vector"]), item["text"])
+        for item in kb
+    ]
     scores.sort(reverse=True, key=lambda x: x[0])
+
     return scores[:top_k], time.time() - t0
 
+
 def build_context(results):
     ctx = ""
     for i, (score, chunk) in enumerate(results):
         ctx += f"=== Context {i+1} ===\n{chunk}\n\n"
     return ctx
 
+
 def rag_answer(query, embed_model, kb):
     t0 = time.time()
+
     results, t_semantic = semantic_search(query, embed_model, kb, top_k=3)
     context = build_context(results)
+
     prompt = f"""Use ONLY the information in the following context.
 
 {context}
@@ -62,19 +123,15 @@ Question: {query}
 If the answer is not in the context, respond EXACTLY with:
 "I do not have enough information to answer that."
 """
-    response = client.chat.completions.create(
-        model="gemini-2.5-pro",
-        temperature=0,
-        messages=[
-            {"role": "system", "content": "Answer strictly using the context."},
-            {"role": "user", "content": prompt}
-        ]
-    )
-    answer = response.choices[0].message.content
+
+    answer = _llm_generate(prompt, max_new_tokens=256)
+
     return answer, t_semantic, time.time() - t0
 
+
 def evaluate_ai(response, true_answer):
     t0 = time.time()
+
     eval_prompt = f"""
 AI Response: {response}
 Ground Truth: {true_answer}
@@ -83,27 +140,35 @@ Rules:
 - 1 = very close to true answer
 - 0.5 = partially correct
 - 0 = incorrect
+
+Answer ONLY one number: 0, 0.5, or 1
 """
-    response = client.chat.completions.create(
-        model="gemini-2.5-pro",
-        temperature=0,
-        messages=[
-            {"role": "system", "content": "You are an evaluation system."},
-            {"role": "user", "content": eval_prompt}
-        ]
-    )
-    return response.choices[0].message.content, time.time() - t0
+
+    score = _llm_generate(eval_prompt, max_new_tokens=8)
+
+    return score.strip(), time.time() - t0
+
 
 def run_rag_pipeline(md_text_input, query, true_answer):
     kb_result = md_to_kb_safe(md_text_input)
+
     if not kb_result["success"]:
         return f"Error creating KB:\n{kb_result['error']}", None, None
+
     kb = kb_result["kb"]
     embed_model = kb_result["embed_model"]
+
     answer, t_semantic, t_rag = rag_answer(query, embed_model, kb)
     score, t_eval = evaluate_ai(answer, true_answer)
-    timings = f"Semantic Search: {t_semantic:.2f}s | LLM Answer: {t_rag:.2f}s | Evaluation: {t_eval:.2f}s"
+
+    timings = (
+        f"Semantic Search: {t_semantic:.2f}s | "
+        f"LLM Answer: {t_rag:.2f}s | "
+        f"Evaluation: {t_eval:.2f}s"
+    )
+
     return answer, score, timings
+
 import base64
 import os
 import re