-
Notifications
You must be signed in to change notification settings - Fork 0
/
neural_searcher.py
40 lines (30 loc) · 1.37 KB
/
neural_searcher.py
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
38
39
40
from qdrant_client import QdrantClient
from sentence_transformers import SentenceTransformer
# sudo docker run -p 6333:6333 \
# -v $(pwd)/qdrant_storage:/qdrant/storage \
# qdrant/qdrant
class NeuralSearcher:
def __init__(self, collection_name):
self.collection_name = collection_name
# Initialize encoder model
self.model = SentenceTransformer("all-MiniLM-L6-v2", device="cpu")
# initialize Qdrant client
self.qdrant_client = QdrantClient("http://localhost:6333")
def search(self, text: str):
# Convert text query into vector
vector = self.model.encode(text).tolist()
# Use `vector` for search for closest vectors in the collection
search_result = self.qdrant_client.search(
collection_name=self.collection_name,
query_vector=vector,
query_filter=None, # If you don't want any filters for now
top=5, # 5 the most closest results is enough
)
# `search_result` contains found vector ids with similarity scores along with the stored payload
# In this function you are interested in payload only
payloads = [hit.payload for hit in search_result]
return payloads
def search_startup(q: str):
return {"result": neural_searcher.search(text=q)}
ns = NeuralSearcher("youtube-search")
print(ns.search("seaside"))