Update app.py

app.py

@@ -25,6 +25,8 @@ from werkzeug.utils import secure_filename
 from geopy.geocoders import Nominatim
 import pickle
 import numpy as np
+import tempfile
+from pathlib import Path
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -50,24 +52,41 @@ model_dir = "./models/llm_model"
 device = "cuda" if torch.cuda.is_available() else "cpu"
 print(f"Using device: {device}")
 
+# Configure cache directories
+os.environ['TRANSFORMERS_CACHE'] = '/cache'
+os.environ['HF_HOME'] = '/cache'
+os.environ['XDG_CACHE_HOME'] = '/cache'
+
 # Load SentenceTransformer model
 def load_sentence_transformer():
     print("Loading SentenceTransformer model...")
     try:
-        model_embedding = SentenceTransformer("jinaai/jina-embeddings-v3", trust_remote_code=True).to(device)
+        # Create cache directory if it doesn't exist
+        cache_dir = Path('/cache')
+        cache_dir.mkdir(parents=True, exist_ok=True)
+        
+        model_embedding = SentenceTransformer(
+            "jinaai/jina-embeddings-v3",
+            trust_remote_code=True,
+            cache_folder=str(cache_dir)
+        ).to(device)
 
         # Load and optimize model state dict
-        state_dict = torch.load(model_path, map_location=device)
-
-        # Dequantize if needed
-        for key, tensor in state_dict.items():
-            if hasattr(tensor, 'dequantize'):  # Check if tensor is quantized
-                state_dict[key] = tensor.dequantize().to(dtype=torch.float32)  # Convert to FP32
-            elif tensor.dtype == torch.bfloat16:  # Handle bfloat16 tensors
-                state_dict[key] = tensor.to(dtype=torch.float32)  # Convert to FP32
-
-        model_embedding.load_state_dict(state_dict)
-        print("SentenceTransformer model loaded successfully.")
+        if os.path.exists(model_path):
+            state_dict = torch.load(model_path, map_location=device)
+
+            # Dequantize if needed
+            for key, tensor in state_dict.items():
+                if hasattr(tensor, 'dequantize'):
+                    state_dict[key] = tensor.dequantize().to(dtype=torch.float32)
+                elif tensor.dtype == torch.bfloat16:
+                    state_dict[key] = tensor.to(dtype=torch.float32)
+
+            model_embedding.load_state_dict(state_dict)
+            print("SentenceTransformer model loaded successfully.")
+        else:
+            print(f"Warning: Model file not found at {model_path}")
+            
         return model_embedding
     except Exception as e:
         print(f"Error loading model: {str(e)}")