QDrantRAG9 / neural_searcher.py
dinhquangson's picture
Create neural_searcher.py
8fd584d verified
raw
history blame
1.46 kB
from qdrant_client import QdrantClient
from sentence_transformers import SentenceTransformer
from qdrant_client.models import Filter
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, city: str):
# Convert text query into vector
vector = self.model.encode(text).tolist()
city_of_interest = city
# Define a filter for cities
city_filter = Filter(**{
"must": [{
"key": "city", # Store city information in a field of the same name
"match": { # This condition checks if payload field has the requested value
"value": city_of_interest
}
}]
})
search_result = self.qdrant_client.search(
collection_name=self.collection_name,
query_vector=vector,
query_filter=city_filter,
limit=5
)
# `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