Vietnamese tokenization (underthesea) + query rewriting

ingest.py

@@ -6,14 +6,17 @@ IMAGE_FOLDER = "data/images"
 if __name__ == "__main__":
     mode = sys.argv[1] if len(sys.argv) > 1 else "all"
 
-    if mode in ("pdf", "all"):
-        import os
-        os.environ["CUDA_VISIBLE_DEVICES"] = ""  # force CPU cho Docling
+    if mode == "all":
+        import subprocess
+        subprocess.run([sys.executable, __file__, "pdf"], check=True)
+        subprocess.run([sys.executable, __file__, "images"], check=True)
+
+    elif mode == "pdf":
         from src.parsers.router import parse_folder as parse_folder_pdf
         from src.chunker import chunk_pages
         from src.vector_store import add_chunks
 
-        print("=== INGEST PDF (Docling - CPU) ===")
+        print("=== INGEST PDF (pypdf - CPU) ===")
         pdf_pages = parse_folder_pdf(PDF_FOLDER, extensions={".pdf"})
         print(f"Tổng: {len(pdf_pages)} đoạn từ PDF")
         chunks = chunk_pages(pdf_pages)
@@ -22,7 +25,7 @@ if __name__ == "__main__":
             add_chunks(chunks)
         print("PDF xong!\n")
 
-    if mode in ("images", "all"):
+    elif mode == "images":
         from src.parsers.router import parse_folder as parse_folder_img
         from src.chunker import chunk_pages
         from src.vector_store import add_chunks
@@ -34,4 +37,4 @@ if __name__ == "__main__":
         print(f"Tổng: {len(chunks)} chunks")
         if chunks:
             add_chunks(chunks)
-        print("Ảnh xong!")
+        print("Ảnh xong!")

src/query_rewriter.py

@@ -0,0 +1,27 @@
+from groq import Groq
+import os
+from dotenv import load_dotenv
+
+load_dotenv()
+_client = Groq(api_key=os.environ["GROQ_API_KEY"])
+
+REWRITE_MODEL = "llama-3.3-70b-versatile"
+
+
+def rewrite_query(query: str) -> str:
+    """Paraphrase câu hỏi để tăng khả năng tìm kiếm."""
+    prompt = f"""Bạn là công cụ cải thiện câu truy vấn tìm kiếm tiếng Việt.
+Hãy viết lại câu hỏi sau thành một câu truy vấn ngắn gọn, súc tích hơn với các từ khóa quan trọng.
+Chỉ trả về câu truy vấn mới, không giải thích.
+
+Câu hỏi gốc: {query}
+Câu truy vấn mới:"""
+
+    response = _client.chat.completions.create(
+        model=REWRITE_MODEL,
+        messages=[{"role": "user", "content": prompt}],
+        max_tokens=100,
+        temperature=0.3,
+    )
+    rewritten = response.choices[0].message.content.strip()
+    return rewritten if rewritten else query

src/rag.py

@@ -4,6 +4,7 @@ from dotenv import load_dotenv
 from src.embedder import embed_query
 from src.vector_store import query as vector_query
 from src.reranker import rerank
+from src.query_rewriter import rewrite_query
 
 load_dotenv()
 client = Groq(api_key=os.environ["GROQ_API_KEY"])
@@ -30,10 +31,13 @@ TRẢ LỜI:"""
 
 
 def answer(question: str, top_k: int = 5) -> dict:
-    """RAG pipeline: embed -> hybrid retrieve -> rerank -> generate."""
-    query_vec = embed_query(question)
-    # Lấy nhiều hơn để reranker có đủ candidates
-    candidates = vector_query(query_vec, query_text=question, top_k=top_k * 3)
+    """RAG pipeline: rewrite -> embed -> hybrid retrieve -> rerank -> generate."""
+    rewritten = rewrite_query(question)
+    if rewritten != question:
+        print(f"Query rewritten: {rewritten}")
+
+    query_vec = embed_query(rewritten)
+    candidates = vector_query(query_vec, query_text=rewritten, top_k=top_k * 3)
     contexts = rerank(question, candidates, top_k=top_k)
 
     prompt = build_prompt(question, contexts)

src/vector_store.py

@@ -5,8 +5,17 @@ from qdrant_client import QdrantClient
 from qdrant_client.models import Distance, VectorParams, PointStruct
 from rank_bm25 import BM25Okapi
 from tqdm import tqdm
+from underthesea import word_tokenize
 from src.embedder import embed_texts
 
+
+def _tokenize_vi(text: str) -> list[str]:
+    """Tách từ tiếng Việt, fallback về split nếu lỗi."""
+    try:
+        return word_tokenize(text.lower(), format="text").split()
+    except Exception:
+        return text.lower().split()
+
 QDRANT_PATH = "qdrant_storage"
 COLLECTION_NAME = "rag_docs"
 BM25_PATH = "bm25_index.pkl"
@@ -53,7 +62,7 @@ def add_chunks(chunks: list[dict]):
 
     # Build BM25 index
     print("Building BM25 index...")
-    tokenized = [t.lower().split() for t in texts]
+    tokenized = [_tokenize_vi(t) for t in texts]
     bm25 = BM25Okapi(tokenized)
     with open(BM25_PATH, "wb") as f:
         pickle.dump({"bm25": bm25, "texts": texts, "metadatas": metadatas}, f)
@@ -95,7 +104,7 @@ def query(query_embedding: list[float], query_text: str, top_k: int = 5) -> list
     texts = bm25_data["texts"]
     metadatas = bm25_data["metadatas"]
 
-    tokenized_query = query_text.lower().split()
+    tokenized_query = _tokenize_vi(query_text)
     bm25_scores = bm25.get_scores(tokenized_query)
     bm25_ids = sorted(range(len(bm25_scores)), key=lambda i: bm25_scores[i], reverse=True)[:top_n]