Update QdrantRag.py

QdrantRag.py

@@ -21,6 +21,10 @@ qdrant_client = QdrantClient(url=os.getenv("qdrant_url"), api_key=os.getenv("qdr
 sparse_encoder = SparseTextEmbedding(model_name="prithivida/Splade_PP_en_v1")
 co = cohere.ClientV2(os.getenv("cohere_api_key"))
 
+dataset = load_dataset("Karbo31881/Pokemon_images")
+ds = dataset["train"]
+labels = ds["text"]
+
 def get_sparse_embedding(text: str, model: SparseTextEmbedding):
     embeddings = list(model.embed(text))
     vector = {f"sparse-text": models.SparseVector(indices=embeddings[0].indices, values=embeddings[0].values)}
@@ -152,7 +156,7 @@ class NeuralSearcher:
         results = co.rerank(model="rerank-v3.5", query=text, documents=search_result, top_n = 3)
         ranked_results = [search_result[results.results[i].index] for i in range(3)]
         return ranked_results
-    def search_image(self, image: ImageFile, limit: int = 1):
+    def search_image(self, image: ImageFile, limit: int = 5):
         img = image
         inputs = self.image_processor(images=img, return_tensors="pt").to(device)
         with torch.no_grad():
@@ -163,5 +167,5 @@ class NeuralSearcher:
             query_filter=None,
             limit=limit,
         )
-        payloads = [hit.payload["label"] for hit in search_result]
-        return payloads
+        payloads = [f"- {hit.payload["label"]} with score {hit.score}" for hit in search_result]
+        return payloads