Update app.py

app.py

@@ -1,3 +1,4 @@
+import os
 import torch
 from loguru import logger
 from pydantic import BaseModel, Field
@@ -30,12 +31,20 @@ logger.warning(
     f"Using device: {DEVICE} ({'GPU: ' + torch.cuda.get_device_name(0) if DEVICE.type == 'cuda' else 'Running on CPU'})"
 )
 
+# Ensure a writable cache directory
+os.makedirs("models_cache", exist_ok=True)
+
 # Load the model at startup to avoid reloading for each request
-model = CrossEncoder(
-    "jinaai/jina-reranker-v1-turbo-en",
-    trust_remote_code=True,
-    device=DEVICE
-)
+try:
+    model = CrossEncoder(
+        "jinaai/jina-reranker-v1-turbo-en",
+        trust_remote_code=True,
+        device=DEVICE,
+        cache_dir="models_cache",
+    )
+except Exception as e:
+    logger.error(f"Failed to load model: {e}")
+    raise RuntimeError("Model loading failed. Check logs for details.")
 
 
 class RerankerRequest(BaseModel):
@@ -71,27 +80,31 @@ async def rerank_documents(request: RerankerRequest):
     - A list of ranked documents with scores and indexes.
     """
     try:
-        # Call the model's rank method with the provided parameters
-        results = model.rank(
-            request.query,
-            request.documents,
-            return_documents=request.return_documents,
-            top_k=request.top_k,
-        )
-
-        # Format the results based on the model's output
+        # Prepare model input
+        inputs = [[request.query, doc] for doc in request.documents]
+
+        # Get ranking scores
+        scores = model.predict(inputs)
+
+        # Sort scores and get top-k results
+        ranked_indices = sorted(
+            range(len(scores)), key=lambda i: scores[i], reverse=True
+        )[: request.top_k]
+
+        # Format results
         formatted_results = [
             RankedResult(
-                score=result["score"],
-                index=result["corpus_id"],
-                document=result["text"] if request.return_documents else None,
+                score=scores[i],
+                index=i,
+                document=request.documents[i] if request.return_documents else None,
             )
-            for result in results
+            for i in ranked_indices
         ]
 
         return RerankerResponse(results=formatted_results)
 
     except Exception as e:
+        logger.error(f"Error in reranking: {e}")
         raise HTTPException(status_code=500, detail=f"Error in reranking: {str(e)}")
 
 
@@ -99,4 +112,4 @@ async def rerank_documents(request: RerankerRequest):
 if __name__ == "__main__":
     import uvicorn
 
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=7860, workers=1)