akashjayampu commited on
Commit
9e62dbe
·
verified ·
1 Parent(s): 4d4009a

Update src/utils.py

Browse files
Files changed (1) hide show
  1. src/utils.py +9 -8
src/utils.py CHANGED
@@ -1,19 +1,20 @@
1
- from sentence_transformers import SentenceTransformer
2
  import faiss
3
  import numpy as np
 
4
 
5
  model = SentenceTransformer("all-MiniLM-L6-v2")
6
 
7
  def get_embeddings(texts):
8
  return model.encode(texts, convert_to_numpy=True)
9
 
10
- def build_index(embeds):
11
- dim = embeds.shape[1]
12
  index = faiss.IndexFlatIP(dim)
13
- faiss.normalize_L2(embeds)
14
- index.add(embeds)
15
  return index
16
 
17
- def search_similar(index, query_emb, k=5):
18
- faiss.normalize_L2(query_emb)
19
- return index.search(query_emb, k)
 
 
 
1
  import faiss
2
  import numpy as np
3
+ from sentence_transformers import SentenceTransformer
4
 
5
  model = SentenceTransformer("all-MiniLM-L6-v2")
6
 
7
  def get_embeddings(texts):
8
  return model.encode(texts, convert_to_numpy=True)
9
 
10
+ def build_index(embeddings):
11
+ dim = embeddings.shape[1]
12
  index = faiss.IndexFlatIP(dim)
13
+ faiss.normalize_L2(embeddings)
14
+ index.add(embeddings)
15
  return index
16
 
17
+ def search_similar(index, query_embedding, k=5):
18
+ faiss.normalize_L2(query_embedding)
19
+ distances, indices = index.search(query_embedding, k)
20
+ return distances, indices