File size: 3,175 Bytes
784ea73 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | #!/usr/bin/env python3
"""Build per-category FAISS indexes (multi-RAG) from UpToDate patient Q&A.
One index per category (30). Each doc = question + answer (English UpToDate).
Query language = Chinese (patients) -> use a multilingual embedding model.
"""
import json
import os
import numpy as np
os.environ.setdefault("HF_ENDPOINT", "https://hf-mirror.com")
ROOT = "/workspace/TASK18"
QA = os.path.join(ROOT, "data/patient_education_qa.jsonl")
OUT = os.path.join(ROOT, "data/rag")
MODEL_NAME = "BAAI/bge-m3" # multilingual (zh query -> en docs)
os.makedirs(OUT, exist_ok=True)
from sentence_transformers import SentenceTransformer
import faiss
model = SentenceTransformer(MODEL_NAME)
model.max_seq_length = 512
print("model loaded:", MODEL_NAME, "| device:", model.device, flush=True)
rows = [json.loads(l) for l in open(QA, encoding="utf-8")]
# group by category
from collections import defaultdict
bycat = defaultdict(list)
for r in rows:
bycat[r["category"]].append(r)
print("categories:", len(bycat), flush=True)
per_cat_stats = {}
for cat, items in bycat.items():
# dedup by question within category
seen = set()
uniq = []
for r in items:
if r["question"] in seen:
continue
seen.add(r["question"])
uniq.append(r)
docs = []
meta = []
for r in uniq:
d = f"{r['question']}\n{r['answer']}"
docs.append(d)
meta.append({"slug": r["slug"], "id": r["id"], "question": r["question"],
"title": r["title"], "level": r["level"], "category": cat})
emb = model.encode(docs, normalize_embeddings=True,
show_progress_bar=True, batch_size=64)
emb = np.asarray(emb, dtype="float32")
idx = faiss.IndexFlatIP(emb.shape[1])
idx.add(emb)
sub = os.path.join(OUT, cat)
os.makedirs(sub, exist_ok=True)
faiss.write_index(idx, os.path.join(sub, "index.faiss"))
with open(os.path.join(sub, "docs.json"), "w", encoding="utf-8") as f:
json.dump(docs, f, ensure_ascii=False)
with open(os.path.join(sub, "meta.json"), "w", encoding="utf-8") as f:
json.dump(meta, f, ensure_ascii=False)
per_cat_stats[cat] = len(uniq)
print(f" [{cat}] docs={len(uniq)}", flush=True)
with open(os.path.join(OUT, "index_map.json"), "w", encoding="utf-8") as f:
json.dump(per_cat_stats, f, ensure_ascii=False, indent=1)
# ---- sanity: multilingual cross-lingual query ----
print("\n=== cross-lingual sanity (Chinese query -> English docs) ===")
q = "我最近2周一直头疼，有什么免吃药的方法缓解"
qe = model.encode([q], normalize_embeddings=True)
for cat in ["brain-and-nerves", "gastrointestinal-system"]:
sub = os.path.join(OUT, cat)
if not os.path.exists(os.path.join(sub, "index.faiss")):
continue
idx = faiss.read_index(os.path.join(sub, "index.faiss"))
meta = json.load(open(os.path.join(sub, "meta.json"), encoding="utf-8"))
docs = json.load(open(os.path.join(sub, "docs.json"), encoding="utf-8"))
D, I = idx.search(qe, 3)
print(f"\n-- category: {cat} --")
for d, i in zip(D[0], I[0]):
print(f" score={d:.3f} | {meta[i]['question'][:70]}")
|