Update main.py

main.py

@@ -15,74 +15,32 @@ CORS(app)
 
 class MyEmbeddingFunction(EmbeddingFunction):
     def embed_documents(self, input: Documents) -> Embeddings:
-        for i in range(5):
-            try:
-                embeddings = []
-                url = "https://api-inference.huggingface.co/models/BAAI/bge-large-en-v1.5"
-                
-                payload = {
-                    "inputs": input
-                }
-                headers = {
-                    'accept': '*/*',
-                    'accept-language': 'en-US,en;q=0.9',
-                    'content-type': 'application/json',
-                    'origin': 'https://huggingface.co',
-                    'priority': 'u=1, i',
-                    'referer': 'https://huggingface.co/',
-                    'sec-ch-ua': '"Google Chrome";v="131", "Chromium";v="131", "Not_A Brand";v="24"',
-                    'sec-ch-ua-mobile': '?0',
-                    'sec-ch-ua-platform': '"Windows"',
-                    'sec-fetch-dest': 'empty',
-                    'sec-fetch-mode': 'cors',
-                    'sec-fetch-site': 'same-site',
-                    'user-agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36',
-                }
-
-                response = requests.post(url, headers=headers, json=payload)
-                return response.json()[0][0]
-            except:
-                pass
+        return self._call_hf_api(input)
 
     def embed_query(self, input: Documents) -> Embeddings:
-        for i in range(5):
-            try:
-                embeddings = []
-                url = "https://api-inference.huggingface.co/models/BAAI/bge-large-en-v1.5"
-
-                payload = {
-                    "inputs": [input]
-                }
-                headers = {
-                    'accept': '*/*',
-                    'accept-language': 'en-US,en;q=0.9',
-                    'content-type': 'application/json',
-                    'origin': 'https://huggingface.co',
-                    'priority': 'u=1, i',
-                    'referer': 'https://huggingface.co/',
-                    'sec-ch-ua': '"Google Chrome";v="131", "Chromium";v="131", "Not_A Brand";v="24"',
-                    'sec-ch-ua-mobile': '?0',
-                    'sec-ch-ua-platform': '"Windows"',
-                    'sec-fetch-dest': 'empty',
-                    'sec-fetch-mode': 'cors',
-                    'sec-fetch-site': 'same-site',
-                    'user-agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36',
-                }
+        return self._call_hf_api([input])
 
-                response = requests.post(url, headers=headers, json=payload)
-                return response.json()[0][0]
-            except Exception as e:
-                print("Error in Embeding :",str(e))
+    def _call_hf_api(self, inputs):
+        url = "https://api-inference.huggingface.co/models/BAAI/bge-large-en-v1.5"
+        headers = {
+            'accept': '*/*',
+            'content-type': 'application/json',
+        }
+        payload = {"inputs": inputs}
+        try:
+            response = requests.post(url, headers=headers, json=payload)
+            response.raise_for_status()
+            return response.json()[0]
+        except Exception as e:
+            print("Embedding API Error:", str(e))
+            return []
 
 try:
     CHROMA_PATH = "chroma"
     custom_embeddings = MyEmbeddingFunction()
-    db = Chroma(
-         persist_directory=CHROMA_PATH,embedding_function=custom_embeddings
-    )
-    #
+    db = Chroma(persist_directory=CHROMA_PATH, embedding_function=custom_embeddings)
 except Exception as e:
-    print("Error in database :",str(e))
+    print("Database Initialization Error:", str(e))
 
 # Initialize the database without persist_directory
 try:
@@ -91,27 +49,20 @@ try:
 
     # Load documents from chroma.sqlite3
     def load_documents_from_sqlite(db_path="chroma.sqlite3"):
-        conn = sqlite3.connect(db_path)
-        cursor = conn.cursor()
-        
-        # Assuming your table structure has "id", "content", and "embedding"
-        cursor.execute("SELECT id, content, embedding FROM documents")
-        rows = cursor.fetchall()
-        
-        collection = db.get_or_create_collection("default_collection")
-        
-        for row in rows:
-            doc_id = row[0]
-            content = row[1]
-            embedding = json.loads(row[2])  # If embeddings are stored as JSON strings
-            collection.add(
-                ids=[doc_id],
-                documents=[content],
-                embeddings=[embedding]
-            )
-        
-        conn.close()
-        print("Loaded documents into Chroma.")
+        try:
+            conn = sqlite3.connect(db_path)
+            cursor = conn.cursor()
+            cursor.execute("SELECT id, content, embedding FROM documents")
+            rows = cursor.fetchall()
+            collection = db.get_or_create_collection("default_collection")
+            for row in rows:
+                doc_id, content, embedding_json = row
+                embedding = json.loads(embedding_json)
+                collection.add(ids=[doc_id], documents=[content], embeddings=[embedding])
+            conn.close()
+            print("Documents loaded into Chroma.")
+        except Exception as e:
+            print("Error loading documents:", str(e))
     
     load_documents_from_sqlite()  # Call to load data
 
@@ -121,44 +72,30 @@ except Exception as e:
 
 def embeddingGen(query):
     url = "https://api-inference.huggingface.co/models/BAAI/bge-large-en-v1.5"
-
-    payload = {
-        "inputs": [query]
-    }
-    headers = {
-        'accept': '*/*',
-        'accept-language': 'en-US,en;q=0.9',
-        'content-type': 'application/json',
-        'origin': 'https://huggingface.co',
-        'priority': 'u=1, i',
-        'referer': 'https://huggingface.co/',
-        'sec-ch-ua': '"Google Chrome";v="131", "Chromium";v="131", "Not_A Brand";v="24"',
-        'sec-ch-ua-mobile': '?0',
-        'sec-ch-ua-platform': '"Windows"',
-        'sec-fetch-dest': 'empty',
-        'sec-fetch-mode': 'cors',
-        'sec-fetch-site': 'same-site',
-        'user-agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36',
-    }
-
-    response = requests.post(url, headers=headers, json=payload)
-    return response.json()[0][0]
+    headers = {'accept': '*/*', 'content-type': 'application/json'}
+    payload = {"inputs": [query]}
+    try:
+        response = requests.post(url, headers=headers, json=payload)
+        response.raise_for_status()
+        return response.json()[0]
+    except Exception as e:
+        print("Embedding Generation Error:", str(e))
+        return []
 
 
 def strings_ranked_by_relatedness(query, df, top_n=5):
-    def relatedness_fn(x, y):
-        x_norm = np.linalg.norm(x)
-        y_norm = np.linalg.norm(y)
-        return np.dot(x, y) / (x_norm * y_norm)
-
-    query_embedding_response = embeddingGen(query)
-    query_embedding = query_embedding_response
-    strings_and_relatednesses = [
-        (row["text"], relatedness_fn(query_embedding, row["embedding"])) for row in df
-    ]
-    strings_and_relatednesses.sort(key=lambda x: x[1], reverse=True)
-    strings, relatednesses = zip(*strings_and_relatednesses)
-    return strings[:top_n], relatednesses[:top_n]
+    def cosine_similarity(x, y):
+        return np.dot(x, y) / (np.linalg.norm(x) * np.linalg.norm(y))
+
+    query_embedding = embeddingGen(query)
+    ranked = sorted(
+        [(row["text"], cosine_similarity(query_embedding, row["embedding"])) for row in df],
+        key=lambda x: x[1],
+        reverse=True
+    )
+    strings, scores = zip(*ranked)
+    return strings[:top_n], scores[:top_n]
+
 
 
 @app.route("/api/gpt", methods=["POST", "GET"])