feat: RAG answers with sutta citations (M1)
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
7
scripts/ask.py
Normal file
7
scripts/ask.py
Normal file
@@ -0,0 +1,7 @@
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from buddhagpt.rag import answer
|
||||
|
||||
r = answer(sys.argv[1], "mlx-community/Qwen2.5-7B-Instruct-4bit", Path("data/lancedb"))
|
||||
print(r["answer"])
|
||||
print("\nSources:", ", ".join(c["uid"] for c in r["citations"]))
|
||||
25
src/buddhagpt/rag.py
Normal file
25
src/buddhagpt/rag.py
Normal file
@@ -0,0 +1,25 @@
|
||||
from pathlib import Path
|
||||
from buddhagpt.index import search
|
||||
|
||||
SYSTEM = (
|
||||
"You are a thoughtful guide grounded in the Pali Canon. Answer with warmth and "
|
||||
"precision. Base doctrinal claims on the provided passages and cite them by uid "
|
||||
"(e.g. [mn21]). If the passages do not cover the question, say so plainly."
|
||||
)
|
||||
|
||||
def build_prompt(question: str, passages: list[dict]) -> list[dict]:
|
||||
ctx = "\n\n".join(f"[{p['uid']}] {p['title']}:\n{p['chunk']}" for p in passages)
|
||||
return [
|
||||
{"role": "system", "content": SYSTEM},
|
||||
{"role": "user", "content": f"Passages:\n{ctx}\n\nQuestion: {question}"},
|
||||
]
|
||||
|
||||
def answer(question: str, model_path: str, db_path: Path) -> dict:
|
||||
from mlx_lm import load, generate
|
||||
passages = search(db_path, question, k=4)
|
||||
model, tokenizer = load(model_path)
|
||||
prompt = tokenizer.apply_chat_template(
|
||||
build_prompt(question, passages), tokenize=False, add_generation_prompt=True
|
||||
)
|
||||
text = generate(model, tokenizer, prompt=prompt, max_tokens=500)
|
||||
return {"answer": text, "citations": [{"uid": p["uid"], "title": p["title"]} for p in passages]}
|
||||
10
tests/test_rag.py
Normal file
10
tests/test_rag.py
Normal file
@@ -0,0 +1,10 @@
|
||||
from buddhagpt.rag import build_prompt
|
||||
|
||||
def test_build_prompt_includes_passages_and_citation_instruction():
|
||||
msgs = build_prompt("What causes suffering?", [
|
||||
{"uid": "sn56.11", "title": "Setting the Wheel in Motion", "chunk": "Craving leads to suffering."}
|
||||
])
|
||||
system = msgs[0]["content"]
|
||||
assert "sn56.11" in msgs[1]["content"]
|
||||
assert "cite" in system.lower()
|
||||
assert msgs[1]["content"].endswith("What causes suffering?")
|
||||
Reference in New Issue
Block a user