Spaces:
Sleeping
Sleeping
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]}") |