Update app.py

app.py

@@ -8,8 +8,14 @@ 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
+import numpy as np
+import requests
+from charset_normalizer import from_bytes
 import zipfile
-import shutil
 import tempfile
 
 # Custom Exception Class
@@ -20,23 +26,27 @@ class GPUQuotaExceededError(Exception):
 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", "/workspace")
 
-# Directories setup
-os.makedirs(PERSISTENT_PATH, exist_ok=True)
+# Persistent storage directories
+PERSISTENT_PATH = os.getenv("PERSISTENT_PATH", "/data")
+os.makedirs(PERSISTENT_PATH, exist_ok=True, mode=0o777)
+
 TEMP_DIR = os.path.join(PERSISTENT_PATH, "temp")
-os.makedirs(TEMP_DIR, exist_ok=True)
+os.makedirs(TEMP_DIR, exist_ok=True, mode=0o777)
+
 OUTPUTS_DIR = os.path.join(PERSISTENT_PATH, "outputs")
-os.makedirs(OUTPUTS_DIR, exist_ok=True)
+os.makedirs(OUTPUTS_DIR, exist_ok=True, mode=0o777)
 
-# Logging Setup
 LOG_DIR = os.getenv("LOG_DIR", os.path.join(PERSISTENT_PATH, "logs"))
-os.makedirs(LOG_DIR, exist_ok=True)
-LOG_FILE = os.path.join(LOG_DIR, "app.log")
+os.makedirs(LOG_DIR, exist_ok=True, mode=0o777)
+
+# Set Hugging Face cache directory to PERSISTENT_PATH
+os.environ["HF_HOME"] = os.path.join(PERSISTENT_PATH, ".huggingface")
+os.makedirs(os.environ["HF_HOME"], exist_ok=True, mode=0o777)
 
+# Logging Setup
 logging.basicConfig(
-    filename=LOG_FILE,
+    filename=os.path.join(LOG_DIR, "app.log"),
     level=logging.INFO,
     format="%(asctime)s - %(levelname)s - %(message)s",
 )
@@ -49,7 +59,7 @@ def initialize_model():
     global model
     try:
         if model is None:
-            model = SentenceTransformer(EMBEDDING_MODEL_NAME, cache_folder=CACHE_DIR)
+            model = SentenceTransformer(EMBEDDING_MODEL_NAME, cache_folder=os.path.join(PERSISTENT_PATH, "models"))
             logger.info(f"Initialized model: {EMBEDDING_MODEL_NAME}")
         return True
     except requests.exceptions.ConnectionError as e:
@@ -132,7 +142,7 @@ def process_files(files):
         all_embeddings = []
         for i in range(0, len(all_chunks), BATCH_SIZE):
             batch = all_chunks[i:i+BATCH_SIZE]
-            embeddings = handle_gpu_operation(lambda: get_model().encode(batch))
+            embeddings = handle_gpu_operation(lambda: get_model().encode(batch)) if model else []
             all_embeddings.extend(embeddings)
 
         # Save results to OUTPUTS_DIR
@@ -157,8 +167,8 @@ def process_files(files):
 @spaces.GPU
 def semantic_search(query, top_k=5):
     global model
-    if model is None:  # Check if model is initialized
-        if not initialize_model():  # Initialize only if needed and within GPU context
+    if model is None:
+        if not initialize_model():
             return "Model initialization failed. Please try again."
 
     try:
@@ -168,10 +178,13 @@ def semantic_search(query, top_k=5):
         # 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
+            chunks = [c for c in chunks if c.strip()]
 
         # Get query embedding
-        query_embedding = handle_gpu_operation(lambda: get_model().encode([query]))[0]
+        if model:
+            query_embedding = handle_gpu_operation(lambda: get_model().encode([query]))[0]
+        else:
+            return "Model not initialized. Please process files first."
 
         # Calculate similarities
         similarities = np.dot(stored_embeddings, query_embedding) / (
@@ -203,7 +216,9 @@ def search_and_format(query, num_results):
 
 def browse_outputs():
     try:
-        os.startfile(OUTPUTS_DIR)  # For Windows, on Linux use subprocess.run(['xdg-open', OUTPUTS_DIR])
+        # Attempt to open the OUTPUTS_DIR
+        os.startfile(OUTPUTS_DIR)
+        return "Opened outputs directory successfully"
     except Exception as e:
         logger.error(f"Error opening file browser: {str(e)}")
         return "Error opening file browser"
@@ -215,16 +230,13 @@ def download_results_from_disk():
             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
+        with tempfile.TemporaryDirectory() as temp_dir:
+            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"
@@ -271,13 +283,13 @@ def create_gradio_interface():
             )
 
             # Download Results Button
-            download_results_button = gr.Button("⬇️ Download Search Results")
+            download_results_button = gr.Button("⬇️ Download Results")
             download_results_button.click(
                 fn=download_results_from_disk,
                 outputs=[gr.File(label="Download Results")]
             )
 
-        with gr.Tab("_FILES_"):
+        with gr.Tab("Outputs"):
             # Browse Outputs Button
             browse_button = gr.Button("📁 Browse Outputs", variant="primary")
             browse_button.click(
@@ -285,13 +297,6 @@ def create_gradio_interface():
                 outputs=None
             )
 
-            # 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(
             process_files,
             inputs=[file_input],