Create rag.py

rag.py

@@ -0,0 +1,136 @@
+# rag.py – retrieval + model generation
+
+import re
+from typing import List, Dict
+
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+from data.loader import ENTRIES, RAW_KNOWLEDGE
+from data.qa_index import answer_from_qa
+
+MODEL_NAME = "SeaLLMs/SeaLLMs-v3-1.5B-Chat"
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME,
+    torch_dtype=torch.float32,  # CPU on HF free tier
+)
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model.to(device)
+
+
+SYSTEM_PROMPT = (
+    "ທ່ານແມ່ນຜູ້ຊ່ວຍເຫຼືອດ້ານປະຫວັດສາດຂອງປະເທດລາວ "
+    "ສໍາລັບນັກຮຽນຊັ້ນ ມ.1. "
+    "ຕອບແຕ່ພາສາລາວ ແລະຕອບສັ້ນໆ 2–3 ປະໂຫຍກ ໃຫ້ເຂົ້າໃຈງ່າຍ. "
+    "ໃຫ້ອີງຈາກຂໍ້ມູນຂ້າງລຸ່ມນີ້ເທົ່ານັ້ນ. "
+    "ຖ້າຂໍ້ມູນບໍ່ພຽງພໍ ຫຼືບໍ່ຊັດເຈນ ໃຫ້ບອກວ່າບໍ່ແນ່ໃຈ."
+)
+
+
+def retrieve_context(question: str, max_entries: int = 2) -> str:
+    """
+    Simple keyword matching over ENTRIES (text + title + keywords).
+    """
+    if not ENTRIES:
+        return RAW_KNOWLEDGE
+
+    q = question.lower().strip()
+    terms = [t for t in re.split(r"\s+", q) if len(t) > 1]
+
+    if not terms:
+        chosen = ENTRIES[:max_entries]
+        return "\n\n".join(
+            f"[ຊັ້ນ {e.get('grade','')}, ບົດ {e.get('chapter','')}, "
+            f"ຫົວຂໍ້ {e.get('section','')} – {e.get('title','')}]\n{e['text']}"
+            for e in chosen
+        )
+
+    scored = []
+
+    for e in ENTRIES:
+        text = e.get("text", "")
+        title = e.get("title", "")
+        kws = e.get("keywords", [])
+        topic = e.get("topic", "")
+
+        base = (text + " " + title).lower()
+        score = 0
+
+        for t in terms:
+            score += base.count(t)
+
+        for kw in kws:
+            kw_lower = kw.lower()
+            for t in terms:
+                if t in kw_lower:
+                    score += 2
+
+        if topic and any(t in topic for t in terms):
+            score += 1
+
+        if score > 0:
+            scored.append((score, e))
+
+    scored.sort(key=lambda x: x[0], reverse=True)
+    top_entries = [e for _, e in scored[:max_entries]]
+
+    if not top_entries:
+        top_entries = ENTRIES[:max_entries]
+
+    blocks = []
+    for e in top_entries:
+        header = (
+            f"[ຊັ້ນ {e.get('grade','')}, "
+            f"ບົດ {e.get('chapter','')}, "
+            f"ຫົວຂໍ້ {e.get('section','')} – {e.get('title','')}]"
+        )
+        blocks.append(f"{header}\n{e.get('text','')}")
+
+    return "\n\n".join(blocks)
+
+
+def build_prompt(question: str) -> str:
+    context = retrieve_context(question)
+    return f"""{SYSTEM_PROMPT}
+
+ຂໍ້ມູນອ້າງອີງ:
+{context}
+
+ຄໍາຖາມ: {question}
+
+ຄໍາຕອບດ້ວຍພາສາລາວ:"""
+
+
+def generate_answer_with_model(question: str) -> str:
+    prompt = build_prompt(question)
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=160,
+            do_sample=False,   # greedy → stable, a bit faster
+        )
+
+    generated_ids = outputs[0][inputs["input_ids"].shape[1]:]
+    answer = tokenizer.decode(generated_ids, skip_special_tokens=True)
+    return answer.strip()
+
+
+def answer_question(question: str) -> str:
+    if not question.strip():
+        return "ກະລຸນາພິມຄໍາຖາມກ່ອນ."
+
+    # 1) try manual + dataset QA first (instant, no model)
+    direct = answer_from_qa(question)
+    if direct:
+        return direct
+
+    # 2) fall back to model + RAG
+    try:
+        return generate_answer_with_model(question)
+    except Exception as e:
+        return f"ລະບົບມີບັນຫາ: {e}"