Update create_faiss.py

create_faiss.py

@@ -3,11 +3,11 @@ import faiss
 from sentence_transformers import SentenceTransformer
 import numpy as np
 
-# Load the Pile dataset (legal text)
-dataset = load_dataset("EleutherAI/the_pile", split="train")
+# Load the US-LegalKit dataset
+dataset = load_dataset("macadeliccc/US-LegalKit", split="train")
 
-# Extract legal-related documents
-law_data = [item['text'] for item in dataset if item['meta']['pile_set_name'] == 'Pile-CC']
+# Extract legal text documents
+law_data = [item['text'] for item in dataset if 'text' in item]
 
 # Load embedding model
 model = SentenceTransformer("all-MiniLM-L6-v2")
@@ -16,10 +16,11 @@ model = SentenceTransformer("all-MiniLM-L6-v2")
 embeddings = model.encode(law_data, convert_to_numpy=True)
 
 # Create FAISS index
-index = faiss.IndexFlatL2(embeddings.shape[1])
-index.add(embeddings)
+dimension = embeddings.shape[1]
+index = faiss.IndexFlatL2(dimension)  # L2 Distance Index
+index.add(embeddings)  # Add vectors to FAISS index
 
 # Save FAISS index
 faiss.write_index(index, "faiss_index.bin")
 
-print("✅ FAISS index saved successfully!")
+print("✅ FAISS index saved successfully as 'faiss_index.bin'!")