Update app.py

app.py

@@ -1,5 +1,8 @@
 import gradio as gr
 import pandas as pd
+import faiss
+import numpy as np
+import os
 from FlagEmbedding import BGEM3FlagModel
 
 # Load the pre-trained embedding model
@@ -11,31 +14,26 @@ df['embeding_context'] = df['embeding_context'].astype(str).fillna('')
 
 # Filter out any rows where 'embeding_context' might be empty or invalid
 df = df[df['embeding_context'] != '']
+    
+index = faiss.read_index('vector_store.index')
+
 
-# Encode the 'embeding_context' column
-embedding_contexts = df['embeding_context'].tolist()
-embeddings_csv = model.encode(embedding_contexts, batch_size=12, max_length=2048)['dense_vecs']
 # Function to perform search and return all columns
 def search_query(query_text):
     num_records = 50
 
     # Encode the input query text
     embeddings_query = model.encode([query_text], batch_size=12, max_length=2048)['dense_vecs']
-    
-    # Compute similarity between the query and the CSV embeddings
-    similarity_matrix = embeddings_query @ embeddings_csv.T
+    embeddings_query_np = np.array(embeddings_query).astype('float32')
 
-    # Rank records by similarity and select the top 'num_records'
-    similarity_scores = similarity_matrix.max(axis=0)
-    top_indices = similarity_scores.argsort()[-num_records:][::-1]
+    # Search in FAISS index for nearest neighbors
+    distances, indices = index.search(embeddings_query_np, num_records)
 
-    # Get the top results and return all columns
-    result_df = df.iloc[top_indices].drop_duplicates(subset=df.columns.difference(['embedding_context']), keep='first')
+    # Get the top results based on FAISS indices
+    result_df = df.iloc[indices[0]].drop(columns=['embeding_context']).drop_duplicates().reset_index(drop=True)
 
-    
     return result_df
 
-
 # Gradio interface function
 def gradio_interface(query_text):
     search_results = search_query(query_text)