Update main.py

main.py

@@ -172,8 +172,12 @@ def fetch_github_repositories(query, max_results=10):
 # ---------------------------
 # Dense Retrieval Model Setup
 # ---------------------------
-model = SentenceTransformer('all-mpnet-base-v2', device='cuda')
-
+try:
+    model = SentenceTransformer('all-mpnet-base-v2', device='cuda')
+except Exception as e:
+    print("Error initializing GPU for SentenceTransformer; falling back to CPU:", e)
+    model = SentenceTransformer('all-mpnet-base-v2', device='cpu')
+    
 def robust_min_max_norm(scores):
     min_val = scores.min()
     max_val = scores.max()
@@ -185,7 +189,11 @@ def robust_min_max_norm(scores):
 # Cross-Encoder Re-Ranking Function
 # ---------------------------
 def cross_encoder_rerank_candidates(candidates, query, model_name, top_n=10):
-    cross_encoder = CrossEncoder(model_name, device='cuda')
+    try:
+        cross_encoder = CrossEncoder(model_name, device='cuda')
+    except Exception as e:
+        print("Error initializing CrossEncoder on GPU; falling back to CPU:", e)
+        cross_encoder = CrossEncoder(model_name, device='cpu')    
     CHUNK_SIZE = 2000
     MAX_DOC_LENGTH = 5000
     MIN_DOC_LENGTH = 200