Update app.py

app.py

@@ -1,30 +1,60 @@
 import os
 import gradio as gr
 import logging
-import numpy as np
-from sentence_transformers import SentenceTransformer
+import traceback
+import spaces
+from typing import Optional, List
+from dataclasses import dataclass
+from datetime import datetime
+from pathlib import Path
+import gc
 import torch
 from torch.amp import autocast
-from spaces import GPU
-import json  # Import json for direct JSON output in UI
+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 tempfile
+import shutil
+
+# Custom Exception Class (Keep this)
+class GPUQuotaExceededError(Exception):
+    pass
 
-# Constants (Keep your HF token secure - use environment variables if possible for real deployments)
+# Constants (Modified Persistent Paths and Cache)
 EMBEDDING_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
-CACHE_DIR = os.getenv("CACHE_DIR", "/tmp/cache")
-PERSISTENT_PATH = os.getenv("PERSISTENT_PATH", "/tmp/data")
-HF_TOKEN = "YOUR_HF_TOKEN"  # REMEMBER TO REPLACE THIS - BEST TO USE ENVIRONMENT VARIABLE
+CHUNK_SIZE = 500
+BATCH_SIZE = 32
+
+# Set Persistent Storage Path (More Explicit Paths - from Worked Code)
+PERSISTENT_PATH = os.getenv("PERSISTENT_PATH", "/data") # Keep this as /data for Spaces persistent storage
+os.makedirs(PERSISTENT_PATH, exist_ok=True, mode=0o777)
 
-# Create directories (still useful to try, even if /tmp/ is ephemeral)
-os.makedirs(CACHE_DIR, exist_ok=True)
-os.makedirs(PERSISTENT_PATH, exist_ok=True)
+# Define Subdirectories (More Explicit Paths)
+TEMP_DIR = os.path.join(PERSISTENT_PATH, "temp")
+os.makedirs(TEMP_DIR, exist_ok=True, mode=0o777)
 
-# Logging Setup (keep logging - it's helpful for debugging)
-LOG_DIR = os.getenv("LOG_DIR", "/data/logs")
-os.makedirs(LOG_DIR, exist_ok=True)
-LOG_FILE = LOG_DIR + "/app.log"
+OUTPUTS_DIR = os.path.join(PERSISTENT_PATH, "outputs")
+os.makedirs(OUTPUTS_DIR, exist_ok=True, mode=0o777)
 
+NPY_CACHE = os.path.join(PERSISTENT_PATH, "npy_cache")
+os.makedirs(NPY_CACHE, exist_ok=True, mode=0o777)
+
+LOG_DIR = os.getenv("LOG_DIR", os.path.join(PERSISTENT_PATH, "logs"))
+os.makedirs(LOG_DIR, exist_ok=True, mode=0o777)
+
+# Set Hugging Face cache directory to persistent storage (From Worked Code - Important!)
+os.environ["HF_HOME"] = os.path.join(PERSISTENT_PATH, ".huggingface")
+os.makedirs(os.environ["HF_HOME"], exist_ok=True, mode=0o777)
+
+# Set Hugging Face token (Keep this - best to use environment variable)
+HF_TOKEN = os.getenv("HF_TOKEN")
+
+# Logging Setup (Keep this - helpful for debugging)
 logging.basicConfig(
-    filename=LOG_FILE,
+    filename=os.path.join(LOG_DIR, "app.log"), # Use os.path.join for log file path
     level=logging.INFO,
     format="%(asctime)s - %(levelname)s - %(message)s",
 )
@@ -32,25 +62,39 @@ logger = logging.getLogger(__name__)
 
 # Model initialization
 model = None
-model_initialization_error = "" # Global variable to store initialization error
+model_initialization_error = "" # Global variable for initialization error
 
 def initialize_model():
+    """
+    Initialize the sentence transformer model with explicit cache path and error handling.
+    Returns:
+        bool: Whether the model was successfully initialized.
+        str: Error message if initialization failed, otherwise empty string.
+    """
     global model, model_initialization_error
     try:
         if model is None:
-            model = SentenceTransformer(EMBEDDING_MODEL_NAME, cache_folder=CACHE_DIR, use_auth_token=HF_TOKEN)
+            model_cache = os.path.join(PERSISTENT_PATH, "models") # Explicit model cache path (from worked code)
+            os.makedirs(model_cache, exist_ok=True, mode=0o777) # Ensure cache directory exists
+            # Use the HF_TOKEN to load the model (as in worked code)
+            model = SentenceTransformer(EMBEDDING_MODEL_NAME, cache_folder=model_cache, use_auth_token=HF_TOKEN)
             logger.info(f"Initialized model: {EMBEDDING_MODEL_NAME}")
             model_initialization_error = "" # Clear any previous error
             return True, "" # Return success and no error message
         return True, "" # Already initialized, return success and no error
-    except Exception as e:
-        error_msg = f"Model initialization failed: {str(e)}"
+    except requests.exceptions.RequestException as e: # Specific network error handling (from worked code)
+        error_msg = f"Connection error during model download: {str(e)}\n{traceback.format_exc()}"
+        logger.error(error_msg)
+        model_initialization_error = error_msg
+        return False, error_msg
+    except Exception as e: # General error handling (from worked code)
+        error_msg = f"Model initialization failed: {str(e)}\n{traceback.format_exc()}"
         logger.error(error_msg)
-        model_initialization_error = error_msg # Store error message
-        return False, error_msg # Return failure and error message
+        model_initialization_error = error_msg
+        return False, error_msg
 
 
-@GPU()
+@spaces.GPU
 def generate_embedding(text, focus):
     global model, model_initialization_error
     if model is None:
@@ -69,11 +113,11 @@ def generate_embedding(text, focus):
         logger.error(error_msg)
         return "", error_msg
 
-@GPU()
+@spaces.GPU
 def save_embedding(embedding_json, name): # Expect JSON string as input from UI
     try:
         embedding = json.loads(embedding_json) # Parse JSON string back to list
-        filepath = f"{PERSISTENT_PATH}/{name}.npy" # Construct full filepath
+        filepath = os.path.join(PERSISTENT_PATH, f"{name}.npy") # Use os.path.join for filepath
         np.save(filepath, np.array(embedding))
         return f"Embedding saved to: {filepath}" # Return filepath in status
     except Exception as e:
@@ -81,10 +125,10 @@ def save_embedding(embedding_json, name): # Expect JSON string as input from UI
         logger.error(error_msg)
         return error_msg
 
-@GPU()
+@spaces.GPU
 def convert_to_json(embedding_json, name): # Expect JSON string as input
     try:
-        filepath = f"{PERSISTENT_PATH}/{name}.json" # Construct full filepath
+        filepath = os.path.join(PERSISTENT_PATH, f"{name}.json") # Use os.path.join for filepath
         with open(filepath, "w") as f:
             f.write(embedding_json) # Directly write the JSON string
         return f"Embedding saved as JSON to: {filepath}" # Return filepath in status
@@ -93,7 +137,7 @@ def convert_to_json(embedding_json, name): # Expect JSON string as input
         logger.error(error_msg)
         return error_msg
 
-@GPU()
+@spaces.GPU
 def process_files(files, focus):
     global model, model_initialization_error
     if model is None:
@@ -106,7 +150,7 @@ def process_files(files, focus):
         file_statuses = [] # To track status for each file
         for file in files:
             try:
-                with open(file.name, 'r') as f:
+                with open(file.name, 'rb') as f:
                     text = f.read()
                 with torch.amp.autocast('cuda'):
                     embedding = model.encode([text])[0].tolist()
@@ -186,7 +230,7 @@ def create_gradio_interface():
         )
 
         download_button.click(
-            lambda name: f"{PERSISTENT_PATH}/{name}.json" if name else None, # Handle empty name
+            lambda name: os.path.join(PERSISTENT_PATH, f"{name}.json") if name else None, # Handle empty name, use os.path.join
             inputs=[save_name_input],
             outputs=[download_output]
         )