File size: 568 Bytes
7eef4b0
 
9e62dbe
7eef4b0
 
 
 
 
 
9e62dbe
 
7eef4b0
9e62dbe
 
7eef4b0
 
9e62dbe
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
import faiss
import numpy as np
from sentence_transformers import SentenceTransformer

model = SentenceTransformer("all-MiniLM-L6-v2")

def get_embeddings(texts):
    return model.encode(texts, convert_to_numpy=True)

def build_index(embeddings):
    dim = embeddings.shape[1]
    index = faiss.IndexFlatIP(dim)
    faiss.normalize_L2(embeddings)
    index.add(embeddings)
    return index

def search_similar(index, query_embedding, k=5):
    faiss.normalize_L2(query_embedding)
    distances, indices = index.search(query_embedding, k)
    return distances, indices