add faiss search

main.py

@@ -1,16 +1,18 @@
 # api/main.py
-from fastapi import FastAPI, HTTPException
 import sentence_transformers
-from huggingface_hub import hf_hub_download, login
 import pandas as pd
+import os
 
+from fastapi import FastAPI, HTTPException
+from huggingface_hub import hf_hub_download, login
 from src.processor import send_to_dataset,search_and_retrieve,generate_tech
 from typing import List, Dict
 from pydantic import BaseModel
-
 from datasets import load_dataset
+from dotenv import load_dotenv
+
+load_dotenv()
 
-import os
 
 login(token=os.getenv("HF_TOKEN"))
 
@@ -23,12 +25,11 @@ app = FastAPI(
 
 
 model = sentence_transformers.SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
-dataset = load_dataset("OrganizedProgrammers/Technologies", streaming=True, split="train")
-
+dataset = load_dataset("OrganizedProgrammers/Technologies", split="train")
+dataset.add_faiss_index(column='embeddings')
 
 class SearchInput(BaseModel):
     title: str
-    type: str = "title"
 
 class SearchOutput(BaseModel):
     title: str
@@ -54,7 +55,7 @@ def post_search(payload: SearchInput):
     """
     Endpoint that returns a search result.
     """
-    config = {"dataset": dataset, "model": model, "type": payload.type}
+    config = {"dataset": dataset, "model": model}
     res = search_and_retrieve(payload.title, config)
     return res
 
@@ -64,7 +65,7 @@ def post_generate_and_push(payload: GenerateInput):
     Endpoint to generate a technology and push it to the dataset
     """
 
-    config = {"dataset": dataset, "model": model, "type": "title"}
+    config = {"dataset": dataset, "model": model}
     res = search_and_retrieve(payload.title, config)
     if res["score"] >= 0.7 and not payload.force:
         raise HTTPException(status_code=500, detail=f"Cannot generate the technology a high score of {res['score']} have been found for the technology : {res['title']}")

requirements.txt

@@ -4,4 +4,6 @@ sentence_transformers
 pandas
 fuzzywuzzy
 google.genai
-datasets
+datasets
+faiss-cpu
+python-dotenv

src/processor.py

@@ -10,33 +10,31 @@ def search_and_retrieve(user_input, config):
     model = config["model"]
 
     user_embedding = model.encode(user_input)
-    results = []
-    max_result = {"score":0, "technology": "", "type":""}
+    results = dataset.get_nearest_examples('embeddings', user_embedding, k=5)
 
-    for row in dataset:
-        name = row["name"]
-        purpose = row["purpose"]
+    s=results.scores
+    t=results.examples
+    n = len(t['name'])
 
-        cosim = model.similarity(row["embeddings"], user_embedding)
-        if config["type"] == "purpose":
-            token_set_ratio = fuzz.token_set_ratio(user_input, purpose)
-        else:
-            token_set_ratio = fuzz.token_set_ratio(user_input, name)
+    result = []
 
-        fuzzy_score = token_set_ratio / 100
-        alpha = 0.6
-        combined_score = alpha * cosim + (1 - alpha) * fuzzy_score
-        result = {"title": name, "purpose": purpose, "score": combined_score.item()}
-        if combined_score > max_result["score"]:
-            max_result = result          
+    for i in range(n):
+        item = {}
 
-        results.append(result)
+        for key, value in t.items():
+          if key!="embeddings":
+            item[key] = value[i]
 
+        result.append(item)
 
-    top_5 = sorted(results, key=lambda x: x['score'], reverse=True)[1:6]
-    max_result["top5"] = top_5
+    for i,r in enumerate(result):
+      r["score"]=float(s[i])
+    
+    final_output = {"title": result[0]["name"], "purpose": result[0]["purpose"], "score": result[0]["score"]}
+    final_output["top5"] = result
+    print(final_output)
 
-    return max_result
+    return final_output
 
 
 def generate_tech(user_input, user_instructions):