./llama-server -m ./hunyuan4b/Q4_K_M.gguf \
  --host 127.0.0.1 --port 8080 \
  --ctx-size 32768 --jinja

from pathlib import Path
import re

def load_docs(folder: str) -> list[str]:
    chunks = []
    for p in sorted(Path(folder).glob("**/*.md")):
        text = p.read_text(encoding="utf-8")
        sections = re.split(r"\n(?=#{1,3} )", text)
        for sec in sections:
            if len(sec) < 200:
                continue
            for i in range(0, len(sec), 1200):
                piece = sec[i:i + 1200]
                if len(piece.strip()) > 50:
                    chunks.append(sec.split("\n")[0] + "\n" + piece)
    return chunks

chunks = load_docs("./docs")
print(len(chunks))

from sentence_transformers import SentenceTransformer

embed_model = SentenceTransformer("BAAI/bge-m3")
vectors = embed_model.encode(chunks, normalize_embeddings=True, batch_size=32)

import numpy as np

query_vec = embed_model.encode([query], normalize_embeddings=True)[0]
scores = np.dot(vectors, query_vec)
top_idx = np.argsort(-scores)[:20]

from sentence_transformers import CrossEncoder
reranker = CrossEncoder("BAAI/bge-reranker-v2-m3")
pairs = [[query, chunks[i]] for i in top_idx]
rerank_scores = reranker.predict(pairs)
final_idx = [top_idx[i] for i in np.argsort(-rerank_scores)[:4]]

from openai import OpenAI

client = OpenAI(base_url="http://127.0.0.1:8080/v1", api_key="local")

def rewrite(query: str) -> str:
    r = client.chat.completions.create(
        model="local-model",
        messages=[
            {"role": "system", "content": "Viết lại câu hỏi của người dùng thành 3 cụm từ khoá dùng cho tìm kiếm vector, mỗi dòng một cụm, không giải thích."},
            {"role": "user", "content": query},
        ],
        temperature=0.1,
    )
    return r.choices[0].message.content

def answer(query: str, context_chunks: list[str]) -> str:
    context = "\n\n---\n\n".join(context_chunks)
    r = client.chat.completions.create(
        model="local-model",
        messages=[
            {"role": "system", "content":
                "Bạn là trợ lý kho tri thức. Chỉ trả lờng dựa trên tài liệu được cung cấp. "
                "Nếu tài liệu không đề cập, hãy nói rõ là chưa có trong tài liệu, đừng suy đoán. "
                "Trả lờng dưới 300 từ và đánh dấu nguồn ở cuối câu bằng [đoạn N]."},
            {"role": "user", "content": f"Tài liệu:\n{context}\n\nCâu hỏi: {query}"},
        ],
        temperature=0.2,
    )
    return r.choices[0].message.content

rewritten = rewrite(query)
print(answer(query, [chunks[i] for i in final_idx]))