Upload folder using huggingface_hub

handler.py

@@ -0,0 +1,62 @@
+from typing import Dict, List, Any
+import torch
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import torch.nn.functional as F
+
+class EndpointHandler:
+    def __init__(self, path: str = "BAAI/bge-reranker-v2-m3"):
+        # Load tokenizer and model
+        self.tokenizer = AutoTokenizer.from_pretrained(path)
+        self.model = AutoModelForSequenceClassification.from_pretrained(path)
+        self.model.eval()
+
+        # Determine the computation device
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.model.to(self.device)
+
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        """
+        Expected input format:
+        {
+            "query": "Your query here",
+            "texts": ["Document 1", "Document 2", ...],
+            "normalize": true  # Optional; defaults to False
+        }
+        """
+        query = data.get("query")
+        texts = data.get("texts", [])
+        normalize = data.get("normalize", False)
+
+        if not query or not texts:
+            return [{"error": "Both 'query' and 'texts' fields are required."}]
+
+        # Prepare input pairs
+        pairs = [[query, text] for text in texts]
+
+        # Tokenize input pairs
+        inputs = self.tokenizer(
+            pairs,
+            padding=True,
+            truncation=True,
+            return_tensors="pt",
+            max_length=512
+        ).to(self.device)
+
+        with torch.no_grad():
+            # Get model logits
+            outputs = self.model(**inputs)
+            scores = outputs.logits.view(-1)
+
+            # Apply sigmoid normalization if requested
+            if normalize:
+                scores = torch.sigmoid(scores)
+
+        # Prepare the response
+        results = [
+            {"index": idx, "score": score.item()}
+            for idx, score in enumerate(scores)
+        ]
+
+        # Sort results by descending score
+        results.sort(key=lambda x: x["score"], reverse=True)
+        return results