summaryrefslogtreecommitdiff
path: root/scripts/run_baseline.py
diff options
context:
space:
mode:
Diffstat (limited to 'scripts/run_baseline.py')
-rw-r--r--scripts/run_baseline.py75
1 files changed, 75 insertions, 0 deletions
diff --git a/scripts/run_baseline.py b/scripts/run_baseline.py
new file mode 100644
index 0000000..74c4710
--- /dev/null
+++ b/scripts/run_baseline.py
@@ -0,0 +1,75 @@
+"""Run vanilla RAG baseline with FAISS retrieval.
+
+Usage:
+ python scripts/run_baseline.py --config configs/default.yaml --memory-bank data/memory_bank.pt --question "Who wrote Hamlet?"
+"""
+
+import argparse
+import logging
+
+import torch
+import yaml
+
+from hag.config import EncoderConfig, GeneratorConfig, HopfieldConfig, PipelineConfig
+from hag.encoder import Encoder
+from hag.generator import Generator
+from hag.memory_bank import MemoryBank
+from hag.retriever_faiss import FAISSRetriever
+from hag.pipeline import RAGPipeline
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+def main() -> None:
+ parser = argparse.ArgumentParser(description="Run vanilla RAG baseline")
+ parser.add_argument("--config", type=str, default="configs/default.yaml")
+ parser.add_argument("--memory-bank", type=str, required=True)
+ parser.add_argument("--question", type=str, required=True)
+ parser.add_argument("--top-k", type=int, default=5)
+ args = parser.parse_args()
+
+ with open(args.config) as f:
+ cfg = yaml.safe_load(f)
+
+ # Override retriever type to faiss
+ pipeline_config = PipelineConfig(
+ hopfield=HopfieldConfig(**{**cfg.get("hopfield", {}), "top_k": args.top_k}),
+ encoder=EncoderConfig(**cfg.get("encoder", {})),
+ generator=GeneratorConfig(**cfg.get("generator", {})),
+ retriever_type="faiss",
+ )
+
+ # Load memory bank to get embeddings for FAISS
+ from hag.config import MemoryBankConfig
+
+ mb = MemoryBank(MemoryBankConfig(**cfg.get("memory", {})))
+ mb.load(args.memory_bank)
+
+ # Build FAISS index from memory bank embeddings
+ import numpy as np
+
+ embeddings_np = mb.embeddings.T.numpy().astype(np.float32) # (N, d)
+ faiss_ret = FAISSRetriever(top_k=args.top_k)
+ faiss_ret.build_index(embeddings_np, mb.passages)
+
+ encoder = Encoder(pipeline_config.encoder)
+ generator = Generator(pipeline_config.generator)
+
+ pipeline = RAGPipeline(
+ config=pipeline_config,
+ encoder=encoder,
+ generator=generator,
+ faiss_retriever=faiss_ret,
+ )
+
+ result = pipeline.run(args.question)
+ print(f"\nQuestion: {result.question}")
+ print(f"Answer: {result.answer}")
+ print(f"\nRetrieved passages:")
+ for i, p in enumerate(result.retrieved_passages):
+ print(f" [{i+1}] {p[:200]}...")
+
+
+if __name__ == "__main__":
+ main()