Update app.py

app.py

@@ -8,14 +8,9 @@ from dataclasses import dataclass
 from datetime import datetime
 from pathlib import Path
 import gc
-
-import torch
-from torch.cuda.amp import autocast
-from transformers import AutoModel, AutoTokenizer
-from sentence_transformers import SentenceTransformer
-from charset_normalizer import from_bytes
-import numpy as np
-import requests
+import zipfile
+import shutil
+import tempfile
 
 # Custom Exception Class
 class GPUQuotaExceededError(Exception):
@@ -26,19 +21,22 @@ EMBEDDING_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
 CHUNK_SIZE = 500
 BATCH_SIZE = 32
 CACHE_DIR = os.getenv("CACHE_DIR", "/tmp/cache")
-PERSISTENT_PATH = os.getenv("PERSISTENT_PATH", "/data")
+PERSISTENT_PATH = os.getenv("PERSISTENT_PATH", "/workspace")
 
-# Create directories
-os.makedirs(CACHE_DIR, exist_ok=True)
+# Directories setup
 os.makedirs(PERSISTENT_PATH, exist_ok=True)
+TEMP_DIR = os.path.join(PERSISTENT_PATH, "temp")
+os.makedirs(TEMP_DIR, exist_ok=True)
+OUTPUTS_DIR = os.path.join(PERSISTENT_PATH, "outputs")
+os.makedirs(OUTPUTS_DIR, exist_ok=True)
 
 # Logging Setup
-LOG_DIR = os.getenv("LOG_DIR", "/data/logs")
+LOG_DIR = os.getenv("LOG_DIR", os.path.join(PERSISTENT_PATH, "logs"))
 os.makedirs(LOG_DIR, exist_ok=True)
-LOG_FILE = Path(LOG_DIR) / "app.log"
+LOG_FILE = os.path.join(LOG_DIR, "app.log")
 
 logging.basicConfig(
-    filename=str(LOG_FILE),
+    filename=LOG_FILE,
     level=logging.INFO,
     format="%(asctime)s - %(levelname)s - %(message)s",
 )
@@ -137,10 +135,12 @@ def process_files(files):
             embeddings = handle_gpu_operation(lambda: get_model().encode(batch))
             all_embeddings.extend(embeddings)
 
-        # Save results
-        np.save(f"{PERSISTENT_PATH}/embeddings.npy", np.array(all_embeddings))
+        # Save results to OUTPUTS_DIR
+        embeddings_path = os.path.join(OUTPUTS_DIR, "embeddings.npy")
+        np.save(embeddings_path, np.array(all_embeddings))
 
-        with open(f"{PERSISTENT_PATH}/chunks.txt", "w", encoding="utf-8") as f:
+        chunks_path = os.path.join(OUTPUTS_DIR, "chunks.txt")
+        with open(chunks_path, "w", encoding="utf-8") as f:
             for chunk in all_chunks:
                 f.write(chunk + "\n===CHUNK_SEPARATOR===\n")
 
@@ -162,16 +162,16 @@ def semantic_search(query, top_k=5):
             return "Model initialization failed. Please try again."
 
     try:
-        # Load saved embeddings
-        stored_embeddings = np.load(f"{PERSISTENT_PATH}/embeddings.npy")
+        # Load saved embeddings from OUTPUTS_DIR
+        stored_embeddings = np.load(os.path.join(OUTPUTS_DIR, "embeddings.npy"))
 
-        # Load stored chunks
-        with open(f"{PERSISTENT_PATH}/chunks.txt", "r", encoding="utf-8") as f:
+        # Load stored chunks from OUTPUTS_DIR
+        with open(os.path.join(OUTPUTS_DIR, "chunks.txt"), "r", encoding="utf-8") as f:
             chunks = f.read().split("\n===CHUNK_SEPARATOR===\n")
             chunks = [c for c in chunks if c.strip()]  # Remove empty chunks
 
         # Get query embedding
-        query_embedding = handle_gpu_operation(lambda: get_model().encode([query]))[0]  # Use get_model() to get the model
+        query_embedding = handle_gpu_operation(lambda: get_model().encode([query]))[0]
 
         # Calculate similarities
         similarities = np.dot(stored_embeddings, query_embedding) / (
@@ -201,40 +201,33 @@ def search_and_format(query, num_results):
         return "Please enter a search query"
     return semantic_search(query, top_k=num_results)
 
-def download_results(text):
-    if not text:
-        return None
-    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-    filename = f"search_results_{timestamp}.txt"
-    with open(filename, "w", encoding="utf-8") as f:
-        f.write(text)
-    return filename
-
-@spaces.GPU
-def safe_generate_embedding(text):
-    global model
-    if model is None:  # Check if model is initialized
-        initialize_model()  # Initialize only if needed and within GPU context
-
+def browse_outputs():
     try:
-        embedding = handle_gpu_operation(
-            lambda: get_model().encode([text])[0].tolist()  # Use get_model() to get the model
-        )
-        return embedding, "", False
-    except GPUQuotaExceededError as e:
-        error_msg = str(e)
-        logger.error(error_msg)
-        return "", error_msg, True
+        os.startfile(OUTPUTS_DIR)  # For Windows, on Linux use subprocess.run(['xdg-open', OUTPUTS_DIR])
     except Exception as e:
-        error_msg = f"Error generating embedding: {str(e)}"
-        logger.error(error_msg)
-        return "", error_msg, True
+        logger.error(f"Error opening file browser: {str(e)}")
+        return "Error opening file browser"
 
-def download_embeddings():
-    embeddings_path = f"{PERSISTENT_PATH}/embeddings.npy"
-    if not os.path.exists(embeddings_path):
-        return None
-    return embeddings_path
+def download_results_from_disk():
+    try:
+        output_files = [
+            os.path.join(OUTPUTS_DIR, "embeddings.npy"),
+            os.path.join(OUTPUTS_DIR, "chunks.txt")
+        ]
+        
+        # Create a temporary zip file
+        temp_dir = tempfile.gettempdir()
+        zip_path = os.path.join(temp_dir, "results.zip")
+        
+        with zipfile.ZipFile(zip_path, 'w') as zipf:
+            for file in output_files:
+                if os.path.exists(file):
+                    zipf.write(file, os.path.basename(file))
+        
+        return zip_path
+    except Exception as e:
+        logger.error(f"Error creating download: {str(e)}")
+        return "Error creating download file"
 
 def create_gradio_interface():
     with gr.Blocks() as demo:
@@ -270,7 +263,6 @@ def create_gradio_interface():
                 lines=10,
                 show_copy_button=True
             )
-            download_button = gr.Button("⬇️ Download Results")
 
             search_button.click(
                 fn=search_and_format,
@@ -278,27 +270,26 @@ def create_gradio_interface():
                 outputs=results_output
             )
 
-            download_button.click(
-                fn=download_results,
-                inputs=[results_output],
-                outputs=[gr.File(label="Download Search Results")]
+            # Download Results Button
+            download_results_button = gr.Button("⬇️ Download Search Results")
+            download_results_button.click(
+                fn=download_results_from_disk,
+                outputs=[gr.File(label="Download Results")]
             )
 
-        with gr.Tab("Inspect Embeddings"):
-            embed_input = gr.Textbox(label="Enter Text for Embedding")
-            embed_button = gr.Button("Generate Embedding")
-            embed_output = gr.Textbox(label="Embedding Vector", lines=5)
-
-            embed_button.click(
-                safe_generate_embedding,
-                inputs=[embed_input],
-                outputs=[embed_output, error_box, error_box]
+        with gr.Tab("_FILES_"):
+            # Browse Outputs Button
+            browse_button = gr.Button("📁 Browse Outputs", variant="primary")
+            browse_button.click(
+                fn=browse_outputs,
+                outputs=None
             )
 
-            download_embeddings_button = gr.Button("⬇️ Download Embeddings")
-            download_embeddings_button.click(
-                fn=download_embeddings,
-                outputs=[gr.File(label="Download Embeddings")]
+            # Download All Results Button
+            download_all_button = gr.Button("⬇️ Download All Results", variant="primary")
+            download_all_button.click(
+                fn=download_results_from_disk,
+                outputs=[gr.File(label="Download All Results")]
             )
 
         process_button.click(