update

app.py

@@ -30,9 +30,14 @@ async def rerank(request: RerankerRequest):
         if not query or not documents:
             raise HTTPException(status_code=400, detail="Query and documents must be provided.")
 
-        # Rerank the documents based on the query
-        pairs = [[query, doc] for doc in documents]
-        scores = model.predict(pairs)
+        from sentence_transformers import util
+
+        # Calculate embeddings for the query and documents
+        query_embedding = model.encode(query, convert_to_tensor=True)
+        document_embeddings = model.encode(documents, convert_to_tensor=True)
+
+        # Calculate cosine similarity between the query and each document
+        scores = util.cos_sim(query_embedding, document_embeddings)[0].tolist()
 
         # Create a list of dictionaries containing the document and its score
         results = [{"document": doc, "score": score} for doc, score in zip(documents, scores)]