File size: 1,550 Bytes
37fbec9 | 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 | """
data/build_faiss_index.py
-------------------------
Builds a FAISS vector database from disease_descriptions.json
using sentence-transformers, laying groundwork for retrieval-augmented explanations.
"""
import json
import faiss
import numpy as np
from pathlib import Path
from sentence_transformers import SentenceTransformer
def build_index(json_path: str, output_index_path: str):
print("Loading disease descriptions...")
with open(json_path, "r") as f:
descriptions = json.load(f)
diseases = list(descriptions.keys())
texts = [descriptions[d] for d in diseases]
print("Loading embedding model (all-MiniLM-L6-v2)...")
model = SentenceTransformer('all-MiniLM-L6-v2')
print("Computing embeddings...")
embeddings = model.encode(texts, show_progress_bar=True)
embeddings = np.array(embeddings).astype("float32")
# Build FAISS Index
dimension = embeddings.shape[1]
index = faiss.IndexFlatL2(dimension)
# Add vectors to index
index.add(embeddings)
print(f"Index built with {index.ntotal} vectors of dimension {dimension}.")
# Save to disk
faiss.write_index(index, output_index_path)
print(f"Saved FAISS index to {output_index_path}")
# Save metadata mapping
metadata_path = Path(output_index_path).with_suffix(".meta.json")
with open(metadata_path, "w") as f:
json.dump(diseases, f)
print(f"Saved metadata to {metadata_path}")
if __name__ == "__main__":
build_index("data/disease_descriptions.json", "data/disease_knowledge.index")
|