File size: 1,153 Bytes
5ecde30 |
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 |
import faiss
import numpy as np
from sentence_transformers import SentenceTransformer
# Load the model and generate embeddings
model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
model_name = 'all-MiniLM-L6-v2'
# Example sentences
sentences = ["This is a test sentence.", "Another example sentence."]
embeddings = model.encode(sentences)
# Convert embeddings to float32
embeddings = np.array(embeddings).astype('float32')
# Create a FAISS index
index = faiss.IndexFlatL2(embeddings.shape[1]) # L2 distance
index.add(embeddings)
# Save the FAISS index
faiss.write_index(index, f"{model_name}_faiss.index")
# Load the FAISS index (for later use)
index = faiss.read_index(f"{model_name}_faiss.index")
# Generate a query embedding
query_sentence = "cat am de platit la factura"
query_embedding = model.encode([query_sentence]).astype('float32')
# Perform similarity search
k = 5 # Number of nearest neighbors to retrieve
D, I = index.search(query_embedding, k) # D: distances, I: indices
# Print results
print(f"Query: {query_sentence}")
print(f"Nearest neighbors indices: {I[0]}")
print(f"Nearest neighbors distances: {D[0]}") |