RAGtest / create_faiss_index.py
willco-afk's picture
Create create_faiss_index.py
d405a7d verified
raw
history blame
633 Bytes
import faiss
import pickle
from sentence_transformers import SentenceTransformer
import chromadb
# Initialize SentenceTransformer model
model = SentenceTransformer('all-MiniLM-L6-v2')
# Load documents or data that you want to index
documents = ['document 1 text', 'document 2 text', 'document 3 text']
# Generate embeddings for documents
embeddings = model.encode(documents)
# Create FAISS index
faiss_index = faiss.IndexFlatL2(embeddings.shape[1]) # Using L2 distance
# Add embeddings to FAISS index
faiss_index.add(embeddings)
# Save the FAISS index
with open('faiss_index.index', 'wb') as f:
pickle.dump(faiss_index, f)