v

app.py

@@ -4,6 +4,7 @@ import faiss
 import numpy as np
 import os
 from FlagEmbedding import BGEM3FlagModel
+from io import BytesIO
 
 # Load the pre-trained embedding model
 model = BGEM3FlagModel('BAAI/bge-m3', use_fp16=True)
@@ -15,36 +16,9 @@ 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'] != '']
 
-# # Encode the 'embeding_context' column
-# embedding_contexts = df['embeding_context'].tolist()
-# embeddings_csv = model.encode(embedding_contexts, batch_size=12, max_length=1024)['dense_vecs']
-
-# # Convert embeddings to numpy array
-# embeddings_np = np.array(embeddings_csv).astype('float32')
-
-# # FAISS index file path
-# index_file_path = 'vector_store_bge_m3.index'
-
-# # Check if FAISS index file already exists
-# if os.path.exists(index_file_path):
-#     # Load the existing FAISS index from file
-#     index = faiss.read_index(index_file_path)
-#     print("FAISS index loaded from file.")
-# else:
-#     # Initialize FAISS index (for L2 similarity)
-#     dim = embeddings_np.shape[1]
-#     index = faiss.IndexFlatL2(dim)
-
-#     # Add embeddings to the FAISS index
-#     index.add(embeddings_np)
-
-#     # Save the FAISS index to a file for future use
-#     faiss.write_index(index, index_file_path)
-#     print("FAISS index created and saved to file.")
-
+# Load the FAISS index
 index = faiss.read_index('vector_store_bge_m3.index')
 
-
 # Function to perform search and return all columns
 def search_query(query_text):
     num_records = 50
@@ -64,7 +38,15 @@ def search_query(query_text):
 # Gradio interface function
 def gradio_interface(query_text):
     search_results = search_query(query_text)
-    return search_results
+    
+    # Save search_results to an Excel file in memory
+    output = BytesIO()
+    with pd.ExcelWriter(output, engine='xlsxwriter') as writer:
+        search_results.to_excel(writer, index=False)
+    excel_data = output.getvalue()
+    
+    # Return the DataFrame and update the download button
+    return search_results, gr.update(value=excel_data)
 
 with gr.Blocks() as app:
     gr.Markdown("<h1>White Stride Red Search (BEG-M3)</h1>")
@@ -78,10 +60,11 @@ with gr.Blocks() as app:
     # Output table for displaying results
     search_output = gr.DataFrame(label="Search Results")
 
-    # Link button click to action
-    search_button.click(fn=gradio_interface, inputs=search_input, outputs=search_output)
-
+    # Download button for Excel file
+    download_button = gr.DownloadButton(label="Download Excel", file_name="search_results.xlsx")
 
+    # Link button click to action
+    search_button.click(fn=gradio_interface, inputs=search_input, outputs=[search_output, download_button])
 
 # Launch the Gradio app
 app.launch()