Create rag_utils.py

rag_utils.py

@@ -0,0 +1,98 @@
+# rag_utils.py
+import os, json
+import numpy as np
+from typing import List, Tuple, Dict
+
+import faiss
+from sentence_transformers import SentenceTransformer
+
+DATA_DIR = os.path.join(os.path.dirname(__file__), "data")
+CACHE_DIR = "/tmp"   # HF Spaces 每次啟動可寫這裡
+
+TAROT_JSON = os.path.join(DATA_DIR, "tarot_data_full.json")
+NUM_JSON   = os.path.join(DATA_DIR, "numerology_data_full.json")
+
+TAROT_IDX  = os.path.join(CACHE_DIR, "faiss_tarot.index")
+TAROT_META = os.path.join(CACHE_DIR, "faiss_tarot_meta.json")
+NUM_IDX    = os.path.join(CACHE_DIR, "faiss_num.index")
+NUM_META   = os.path.join(CACHE_DIR, "faiss_num_meta.json")
+
+# 輕量好用
+EMBED_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
+
+_model = None
+def get_model():
+    global _model
+    if _model is None:
+        _model = SentenceTransformer(EMBED_MODEL_NAME)
+    return _model
+
+def _build_tarot():
+    with open(TAROT_JSON) as f:
+        data = json.load(f)
+
+    texts = [
+        (i, c["card_name"],
+         (c.get("meaning_upright","") + " " + c.get("advice","")).strip())
+        for i, c in enumerate(data)
+    ]
+    model = get_model()
+    embs = model.encode([t[2] for t in texts], normalize_embeddings=True)
+
+    idx = faiss.IndexFlatIP(embs.shape[1])
+    idx.add(np.array(embs, dtype="float32"))
+
+    faiss.write_index(idx, TAROT_IDX)
+    with open(TAROT_META, "w") as f:
+        json.dump([{"i":i, "card_name":name, "text":txt} for (i,name,txt) in texts], f, indent=2)
+
+def _build_num():
+    with open(NUM_JSON) as f:
+        data = json.load(f)
+
+    texts = [
+        (i, n["number"],
+         (str(n["number"]) + " " + n.get("life_path_meaning","") + " " + n.get("advice","")).strip())
+        for i, n in enumerate(data)
+    ]
+    model = get_model()
+    embs = model.encode([t[2] for t in texts], normalize_embeddings=True)
+
+    idx = faiss.IndexFlatIP(embs.shape[1])
+    idx.add(np.array(embs, dtype="float32"))
+
+    faiss.write_index(idx, NUM_IDX)
+    with open(NUM_META, "w") as f:
+        json.dump([{"i":i, "number":num, "text":txt} for (i,num,txt) in texts], f, indent=2)
+
+def ensure_indexes():
+    os.makedirs(CACHE_DIR, exist_ok=True)
+    if not (os.path.exists(TAROT_IDX) and os.path.exists(TAROT_META)):
+        _build_tarot()
+    if not (os.path.exists(NUM_IDX) and os.path.exists(NUM_META)):
+        _build_num()
+
+def _search(index_path: str, meta_path: str, query: str, k: int = 3):
+    model = get_model()
+    idx = faiss.read_index(index_path)
+    with open(meta_path) as f:
+        meta = json.load(f)
+
+    q = model.encode([query], normalize_embeddings=True).astype("float32")
+    D, I = idx.search(q, k)
+    results = []
+    for rank, (score, j) in enumerate(zip(D[0], I[0]), start=1):
+        m = meta[j]
+        m = dict(m)  # copy
+        m["score"] = float(score)
+        m["rank"] = rank
+        results.append(m)
+    return results
+
+def search_tarot(query: str, k: int = 3):
+    ensure_indexes()
+    return _search(TAROT_IDX, TAROT_META, query, k)
+
+def search_numerology(query: str, k: int = 3):
+    ensure_indexes()
+    return _search(NUM_IDX, NUM_META, query, k)