Update app.py

app.py

@@ -2,47 +2,68 @@ import gradio as gr
 import pandas as pd
 import numpy as np
 from sentence_transformers import SentenceTransformer
-from sklearn.metrics.pairwise import cosine_similarity
+import chromadb
+from chromadb.config import Settings
+from chromadb.utils import embedding_functions
 
-# Load the model
-model = SentenceTransformer('all-MiniLM-L6-v2')
+# Constants
+CSV_PATH = "path/to/your/csv/file.csv"  # Replace with your actual CSV file path
+COLLECTION_NAME = "query_embeddings"
+EMBEDDING_MODEL = 'all-MiniLM-L6-v2'
 
-def load_csv(file):
-    df = pd.read_csv(file.name)
-    return df
+# Initialize embedding model
+embedding_function = embedding_functions.SentenceTransformerEmbeddingFunction(model_name=EMBEDDING_MODEL)
 
-def search_similar_queries(query, df, top_k=5):
-    # Encode the query
-    query_embedding = model.encode([query])
+# Initialize Chroma client
+chroma_client = chromadb.Client(Settings(persist_directory="./chroma_db"))
 
-    # Encode all queries in the DataFrame
-    all_embeddings = model.encode(df['query'].tolist())
-
-    # Calculate cosine similarity
-    similarities = cosine_similarity(query_embedding, all_embeddings)[0]
-
-    # Get top-k similar queries
-    top_indices = np.argsort(similarities)[-top_k:][::-1]
+def load_and_embed_data():
+    # Load CSV
+    df = pd.read_csv(CSV_PATH)
+    
+    # Create or get collection
+    collection = chroma_client.get_or_create_collection(name=COLLECTION_NAME, embedding_function=embedding_function)
+    
+    # Check if collection is empty
+    if collection.count() == 0:
+        print("Embedding and storing data. This may take a while...")
+        # Embed and store data
+        collection.add(
+            documents=df['query'].tolist(),
+            metadatas=df.to_dict('records'),
+            ids=[str(i) for i in range(len(df))]
+        )
+        print("Data embedded and stored.")
+    else:
+        print("Data already embedded and stored.")
+    
+    return collection
 
-    results = []
-    for idx in top_indices:
+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': df.iloc[idx]['query'],
-            'similarity': similarities[idx],
-            'uber_intent': df.iloc[idx]['uber_intent'],
-            'common_intent': df.iloc[idx]['common_intent'],
-            'sub_common_intent': df.iloc[idx]['sub_common_intent'],
-            'fsc': df.iloc[idx]['fsc'],
-            'language': df.iloc[idx]['language'],
-            'Name': df.iloc[idx]['Name']
+            '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']
         }
-        results.append(result)
-
-    return results
+        formatted_results.append(result)
+    
+    return formatted_results
 
-def gradio_interface(csv_file, query, top_k):
-    df = load_csv(csv_file)
-    results = search_similar_queries(query, df, top_k)
+def gradio_interface(query, top_k):
+    results = search_similar_queries(query, collection, top_k)
     
     output = ""
     for i, result in enumerate(results, 1):
@@ -58,17 +79,19 @@ def gradio_interface(csv_file, query, top_k):
     
     return output
 
+# Load and embed data
+collection = load_and_embed_data()
+
 # Create Gradio interface
 iface = gr.Interface(
     fn=gradio_interface,
     inputs=[
-        gr.File(label="Upload CSV file"),
         gr.Textbox(label="Enter your query"),
         gr.Slider(minimum=1, maximum=10, step=1, label="Top-K results", value=5)
     ],
     outputs=gr.Textbox(label="Results"),
-    title="Query Similarity Search",
-    description="Upload a CSV file, enter a query, and find similar queries with associated metadata."
+    title="Optimized Query Similarity Search",
+    description="Enter a query to find similar queries with associated metadata."
 )
 
 # Launch the interface