Update app.py

app.py

@@ -1,7 +1,5 @@
 import gradio as gr
 import pandas as pd
-import numpy as np
-from sentence_transformers import SentenceTransformer
 import chromadb
 from chromadb.config import Settings
 from chromadb.utils import embedding_functions
@@ -11,7 +9,7 @@ CSV_PATH = "ime-cdata - Sheet1.csv"  # Replace with your actual CSV file path
 COLLECTION_NAME = "query_embeddings"
 EMBEDDING_MODEL = 'all-MiniLM-L6-v2'
 
-# Initialize embedding model
+# Initialize embedding function
 embedding_function = embedding_functions.SentenceTransformerEmbeddingFunction(model_name=EMBEDDING_MODEL)
 
 # Initialize Chroma client
@@ -27,7 +25,7 @@ def load_and_embed_data():
     # Check if collection is empty
     if collection.count() == 0:
         print("Embedding and storing data. This may take a while...")
-        # Embed and store data
+        # Embed and store only the query column
         collection.add(
             documents=df['query'].tolist(),
             metadatas=df.to_dict('records'),
@@ -40,44 +38,56 @@ def load_and_embed_data():
     return collection
 
 def search_similar_queries(query, collection, top_k=5):
-    results = collection.query(
-        query_texts=[query],
-        n_results=top_k,
-        include=["metadatas", "distances"]
-    )
-    
-    formatted_results = []
-    for i in range(len(results['ids'][0])):
-        result = {
-            'query': results['metadatas'][0][i]['query'],
-            'similarity': 1 - results['distances'][0][i],  # Convert distance to similarity
-            'uber_intent': results['metadatas'][0][i]['uber_intent'],
-            'common_intent': results['metadatas'][0][i]['common_intent'],
-            'sub_common_intent': results['metadatas'][0][i]['sub_common_intent'],
-            'fsc': results['metadatas'][0][i]['fsc'],
-            'language': results['metadatas'][0][i]['language'],
-            'Name': results['metadatas'][0][i]['Name']
-        }
-        formatted_results.append(result)
-    
-    return formatted_results
+    try:
+        results = collection.query(
+            query_texts=[query],
+            n_results=top_k,
+            include=["metadatas", "distances"]
+        )
+        
+        formatted_results = []
+        for i in range(len(results['ids'][0])):
+            metadata = results['metadatas'][0][i]
+            result = {
+                'query': metadata.get('query', 'N/A'),
+                'similarity': 1 - results['distances'][0][i],  # Convert distance to similarity
+                'uber_intent': metadata.get('uber_intent', 'N/A'),
+                'common_intent': metadata.get('common_intent', 'N/A'),
+                'sub_common_intent': metadata.get('sub_common_intent', 'N/A'),
+                'fsc': metadata.get('fsc', 'N/A'),
+                'language': metadata.get('language', 'N/A'),
+                'Name': metadata.get('Name', 'N/A')
+            }
+            formatted_results.append(result)
+        
+        return formatted_results
+    except Exception as e:
+        print(f"Error in search_similar_queries: {str(e)}")
+        return []
 
 def gradio_interface(query, top_k):
-    results = search_similar_queries(query, collection, top_k)
-    
-    output = ""
-    for i, result in enumerate(results, 1):
-        output += f"Result {i}:\n"
-        output += f"Query: {result['query']}\n"
-        output += f"Similarity: {result['similarity']:.4f}\n"
-        output += f"Uber Intent: {result['uber_intent']}\n"
-        output += f"Common Intent: {result['common_intent']}\n"
-        output += f"Sub-Common Intent: {result['sub_common_intent']}\n"
-        output += f"FSC: {result['fsc']}\n"
-        output += f"Language: {result['language']}\n"
-        output += f"Name: {result['Name']}\n\n"
-    
-    return output
+    try:
+        results = search_similar_queries(query, collection, top_k)
+        
+        if not results:
+            return "No results found or an error occurred."
+        
+        output = ""
+        for i, result in enumerate(results, 1):
+            output += f"Result {i}:\n"
+            output += f"Query: {result['query']}\n"
+            output += f"Similarity: {result['similarity']:.4f}\n"
+            output += f"Uber Intent: {result['uber_intent']}\n"
+            output += f"Common Intent: {result['common_intent']}\n"
+            output += f"Sub-Common Intent: {result['sub_common_intent']}\n"
+            output += f"FSC: {result['fsc']}\n"
+            output += f"Language: {result['language']}\n"
+            output += f"Name: {result['Name']}\n\n"
+        
+        return output
+    except Exception as e:
+        print(f"Error in gradio_interface: {str(e)}")
+        return f"An error occurred: {str(e)}"
 
 # Load and embed data
 collection = load_and_embed_data()
@@ -90,7 +100,7 @@ iface = gr.Interface(
         gr.Slider(minimum=1, maximum=10, step=1, label="Top-K results", value=5)
     ],
     outputs=gr.Textbox(label="Results"),
-    title="Optimized Query Similarity Search",
+    title="Query Similarity Search",
     description="Enter a query to find similar queries with associated metadata."
 )