Add prebuilt FAISS index and loader for Hugging Face Space

build_faiss_index.py

@@ -0,0 +1,39 @@
+# build_faiss_index.py
+import faiss
+import pickle
+import numpy as np
+from datasets import load_dataset
+from langchain.embeddings import HuggingFaceEmbeddings  # Or your embedding class
+
+def format_entry(entry):
+    return {
+        "question": entry["question"],
+        "formatted": (
+            f"Q: {entry['question']}\n"
+            f"A. {entry['opa']}  B. {entry['opb']}  C. {entry['opc']}  D. {entry['opd']}\n"
+            f"Correct Answer: {entry['cop']}\n"
+            f"Explanation: {entry['exp']}"
+        )
+    }
+
+print("Loading MedMCQA dataset (5000 rows)...")
+dataset = load_dataset("medmcqa", split="train[:5000]")
+formatted_data = [format_entry(entry) for entry in dataset]
+
+# Extract questions for embeddings
+questions = [entry["formatted"] for entry in formatted_data]
+
+print("Generating embeddings...")
+embeddings_model = HuggingFaceEmbeddings()  # Or your specific embeddings
+vectors = np.array([embeddings_model.embed_query(q) for q in questions], dtype="float32")
+
+print("Building FAISS index...")
+index = faiss.IndexFlatL2(vectors.shape[1])
+index.add(vectors)
+
+# Save FAISS index and data
+faiss.write_index(index, "data/medmcqa_index/index.faiss")
+with open("data/medmcqa_index/index.pkl", "wb") as f:
+    pickle.dump(formatted_data, f)
+
+print("FAISS index saved successfully!")

dataset_loader.py

@@ -1,18 +1,11 @@
 # dataset_loader.py
-from datasets import load_dataset
+import pickle
+import faiss
 
-def load_medmcqa_subset(limit=5000):
-    dataset = load_dataset("medmcqa", split=f"train[:{limit}]")
-
-    def format_entry(entry):
-        return {
-            "question": entry["question"],
-            "formatted": (
-                f"Q: {entry['question']}\n"
-                f"A. {entry['opa']}  B. {entry['opb']}  C. {entry['opc']}  D. {entry['opd']}\n"
-                f"Correct Answer: {entry['cop']}\n"
-                f"Explanation: {entry['exp']}"
-            )
-        }
-
-    return [format_entry(entry) for entry in dataset]
+def load_faiss_index():
+    print("Loading FAISS index...")
+    index = faiss.read_index("data/medmcqa_index/index.faiss")
+    with open("data/medmcqa_index/index.pkl", "rb") as f:
+        formatted_data = pickle.load(f)
+    print("FAISS index loaded successfully!")
+    return index, formatted_data