Update app.py

app.py

@@ -14,6 +14,7 @@ from pathlib import Path
 import tempfile
 import spaces  # for ZeroGPU
 import requests  # for IP geolocation
+import time
 
 # Load environment variables and initialize clients
 load_dotenv()
@@ -98,23 +99,51 @@ def initialize_system_sync():
     # Use the same ChromaDB client that was loaded from HF
     chroma_client = db  # Use the global db instance we created
     
-    # Initialize the embedding function
-    embedding_function = chromadb.utils.embedding_functions.SentenceTransformerEmbeddingFunction(
-        model_name="sentence-transformers/all-mpnet-base-v2",
-        device=DEVICE
-    )
+    # Initialize the embedding function with retries
+    max_retries = 3
+    retry_delay = 5  # seconds
+    
+    for attempt in range(max_retries):
+        try:
+            print(f"\nAttempt {attempt + 1} of {max_retries} to initialize embedding function...")
+            embedding_function = chromadb.utils.embedding_functions.SentenceTransformerEmbeddingFunction(
+                model_name="sentence-transformers/all-mpnet-base-v2",
+                device=DEVICE
+            )
+            break
+        except Exception as e:
+            print(f"Error initializing embedding function: {str(e)}")
+            if attempt < max_retries - 1:
+                print(f"Retrying in {retry_delay} seconds...")
+                time.sleep(retry_delay)
+            else:
+                raise RuntimeError("Failed to initialize embedding function after multiple attempts")
     
     # Get the collection
     print("Getting collection...")
     collection = chroma_client.get_collection(name="website_content", embedding_function=embedding_function)
     print(f"Found {collection.count()} documents in collection")
     
-    # Initialize the reranker and explicitly move to GPU if available
-    print("\nInitialising Cross-Encoder...")
-    reranker = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2', device=DEVICE)
-    if torch.cuda.is_available():
-        reranker.model.to('cuda')  # Ensure model is on GPU
-        print("Reranker moved to GPU")
+    # Initialize the reranker with retries
+    for attempt in range(max_retries):
+        try:
+            print(f"\nAttempt {attempt + 1} of {max_retries} to initialize reranker...")
+            reranker = CrossEncoder(
+                'cross-encoder/ms-marco-MiniLM-L-6-v2',
+                device=DEVICE,
+                max_length=512  # Add explicit max_length
+            )
+            if torch.cuda.is_available():
+                reranker.model.to('cuda')
+                print("Reranker moved to GPU")
+            break
+        except Exception as e:
+            print(f"Error initializing reranker: {str(e)}")
+            if attempt < max_retries - 1:
+                print(f"Retrying in {retry_delay} seconds...")
+                time.sleep(retry_delay)
+            else:
+                raise RuntimeError("Failed to initialize reranker after multiple attempts")
 
 @spaces.GPU(memory="40g")
 def initialize_system():