./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": "把用户问题改写成 3 个用于向量检索的关键词短语，用换行分隔，不要解释。"},
            {"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":
                "你是知识库助手。只依据给定材料回答问题。"
                "材料中没有的信息，明确回答“材料未提及”，不要推测。"
                "回答控制在 300 字以内，并在句末用 [片段N] 标注依据。"},
            {"role": "user", "content": f"材料：\n{context}\n\n问题：{query}"},
        ],
        temperature=0.2,
    )
    return r.choices[0].message.content

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