Upload examples/run_memory.py with huggingface_hub

examples/run_memory.py

@@ -1,12 +1,13 @@
 """Matrix-BIOS-Memory-0.1 — grounded, citation-faithful recall (RAG).
-Ships a FAISS index + a small corpus; answers cite the source ids they used.
+Ships a FAISS index + a small corpus; every answer cites the source ids it used.
 pip install torch transformers sentence-transformers faiss-cpu huggingface_hub
 """
 import json
 import faiss
+import torch
 from huggingface_hub import snapshot_download
 from sentence_transformers import SentenceTransformer
-from transformers import pipeline
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
 REPO = "ruslanmv/Matrix-BIOS-Memory-0.1"
 path = snapshot_download(REPO)
@@ -14,8 +15,9 @@ cfg  = json.load(open(f"{path}/memory_config.json"))   # embedder / generator /
 docs = json.load(open(f"{path}/docs.json"))            # [{"id": ..., "text": ...}]
 index = faiss.read_index(f"{path}/index.faiss")
 
-embedder  = SentenceTransformer(cfg["embedder"])
-generator = pipeline("text2text-generation", model=cfg["generator"])
+embedder = SentenceTransformer(cfg["embedder"])
+gen_tok  = AutoTokenizer.from_pretrained(cfg["generator"])
+gen_model = AutoModelForSeq2SeqLM.from_pretrained(cfg["generator"]).eval()
 
 def answer(question: str):
     qv = embedder.encode([question], normalize_embeddings=True).astype("float32")
@@ -24,8 +26,10 @@ def answer(question: str):
     context = "\n".join(f"[{d['id']}] {d['text']}" for d in hits)
     prompt = ("Answer the question using ONLY the context, and cite the [id] you used.\n"
               f"Context:\n{context}\n\nQuestion: {question}\nAnswer:")
-    out = generator(prompt, max_new_tokens=64)[0]["generated_text"]
-    return out, [d["id"] for d in hits]
+    ids = gen_tok(prompt, return_tensors="pt", truncation=True).input_ids
+    with torch.no_grad():
+        out = gen_model.generate(ids, max_new_tokens=64)
+    return gen_tok.decode(out[0], skip_special_tokens=True), [d["id"] for d in hits]
 
 if __name__ == "__main__":
     for q in ["What does every effectful action in Matrix OS emit?",