Update app.py

app.py

@@ -17,7 +17,7 @@ import requests
 from charset_normalizer import from_bytes
 import zipfile
 import tempfile
-import webbrowser
+import shutil
 
 # Custom Exception Class
 class GPUQuotaExceededError(Exception):
@@ -125,15 +125,15 @@ def get_model():
 @spaces.GPU
 def process_files(files):
     if not files:
-        return "Please upload one or more.txt files.", "", ""
+        return "Please upload one or more.txt files.", ""
 
     try:
         if not initialize_model():
-            return "Failed to initialize the model. Please try again.", "", ""
+            return "Failed to initialize the model. Please try again.", ""
 
         valid_files = [f for f in files if f.name.lower().endswith('.txt')]
         if not valid_files:
-            return "No.txt files found. Please upload valid.txt files.", "", ""
+            return "No.txt files found. Please upload valid.txt files.", ""
 
         all_chunks = []
         processed_files = 0
@@ -154,7 +154,7 @@ def process_files(files):
                 logger.error(f"Error processing file {file.name}: {str(e)}")
 
         if not all_chunks:
-            return "No valid content found in the uploaded files.", "", ""
+            return "No valid content found in the uploaded files.", ""
 
         # Generate embeddings in batches
         all_embeddings = []
@@ -164,7 +164,7 @@ def process_files(files):
                 embeddings = handle_gpu_operation(lambda: model.encode(batch))
                 all_embeddings.extend(embeddings)
             else:
-                return "Model not initialized. Please check model initialization.", "", ""
+                return "Model not initialized. Please check model initialization.", ""
 
         # Save results to OUTPUTS_DIR
         embeddings_path = os.path.join(OUTPUTS_DIR, "embeddings.npy")
@@ -176,13 +176,12 @@ def process_files(files):
 
         return (
             f"Successfully processed {processed_files} files. Generated {len(all_embeddings)} embeddings from {len(all_chunks)} chunks.",
-            "",
             ""
         )
 
     except Exception as e:
         logger.error(f"Processing failed: {str(e)}")
-        return f"Error processing files: {str(e)}", "", ""
+        return f"Error processing files: {str(e)}", ""
 
 @spaces.GPU
 def semantic_search(query, top_k=5):
@@ -194,13 +193,6 @@ def semantic_search(query, top_k=5):
         # Load saved embeddings and chunks from OUTPUTS_DIR
         embeddings_file = os.path.join(OUTPUTS_DIR, "embeddings.npy")
         chunks_file = os.path.join(OUTPUTS_DIR, "chunks.txt")
-
-        logger.info(f"Checking for embeddings file: {embeddings_file}")
-        logger.info(f"Checking for chunks file: {chunks_file}")
-
-        if not os.path.exists(embeddings_file) or not os.path.exists(chunks_file):
-            return "Embeddings or chunks not found. Please generate embeddings first."
-
         stored_embeddings = np.load(embeddings_file)
         with open(chunks_file, "r", encoding="utf-8") as f:
             chunks = f.read().split("\n===CHUNK_SEPARATOR===\n")
@@ -233,31 +225,17 @@ def search_and_format(query, num_results):
         return "Please enter a search query"
     return semantic_search(query, top_k=num_results)
 
-def browse_outputs():
+def copy_embeddings_to_workspace():
     try:
-        # Open the outputs directory in a web browser (may work on some systems)
-        webbrowser.open(f"file://{OUTPUTS_DIR}")
-        return "Opened outputs directory."
-    except Exception as e:
-        logger.error(f"Error opening file browser: {str(e)}")
-        return "Error opening file browser."
-
-def download_results():
-    required_files = ["embeddings.npy", "chunks.txt"]
-    missing = [f for f in required_files if not os.path.exists(os.path.join(OUTPUTS_DIR, f))]
-    if missing:
-        logger.error(f"Missing files: {missing}")
-        return None
-    try:
-        zip_path = os.path.join(OUTPUTS_DIR, "results.zip")
-        with zipfile.ZipFile(zip_path, 'w') as zipf:
-            for file in required_files:
-                file_path = os.path.join(OUTPUTS_DIR, file)
-                zipf.write(file_path, file)
-        return zip_path
+        embeddings_path = os.path.join(OUTPUTS_DIR, "embeddings.npy")
+        chunks_path = os.path.join(OUTPUTS_DIR, "chunks.txt")
+        workspace_dir = os.getcwd()
+        shutil.copy(embeddings_path, workspace_dir)
+        shutil.copy(chunks_path, workspace_dir)
+        return "Embeddings copied to workspace directory."
     except Exception as e:
-        logger.error(f"Error creating download archive: {str(e)}")
-        return None
+        logger.error(f"Error copying embeddings: {str(e)}")
+        return f"Error copying embeddings: {str(e)}"
 
 def create_gradio_interface():
     with gr.Blocks() as demo:
@@ -278,7 +256,7 @@ def create_gradio_interface():
         process_button.click(
             fn=process_files,
             inputs=[file_input],
-            outputs=[output_text, error_box, error_box]
+            outputs=[output_text, error_box]
         )
 
         with gr.Tab("Search"):
@@ -305,18 +283,12 @@ def create_gradio_interface():
                 outputs=results_output
             )
 
-            download_button = gr.Button(" Download Results")
-            download_button.click(
-                fn=download_results,
-                outputs=[gr.File(label="Download Results")]
-            )
-
-        with gr.Tab("Outputs"):
-            browse_button = gr.Button(" Browse Outputs")
-            browse_button.click(
-                fn=browse_outputs,
-                outputs=[gr.Textbox(label="Browse Status")]
-            )
+        copy_button = gr.Button("Copy Embeddings to Workspace")
+        copy_output = gr.Textbox(label="Copy Status")
+        copy_button.click(
+            fn=copy_embeddings_to_workspace,
+            outputs=[copy_output]
+        )
 
     return demo