open-paws/text_performance_prediction_longform

1from transformers import AutoTokenizer, AutoModelForSequenceClassification2import torch3 4class EndpointHandler:5    def __init__(self, path=""):6        # Load model and tokenizer from the repo path7        self.tokenizer = AutoTokenizer.from_pretrained(path)8        self.model = AutoModelForSequenceClassification.from_pretrained(path)9        self.model.eval()10        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")11        self.model.to(self.device)12 13    def __call__(self, data):14        """15        This method is called when the endpoint receives a request.16        Expected input: { "inputs": "some string" } or { "inputs": ["a", "b", ...] }17        """18        inputs = data.get("inputs", None)19 20        if inputs is None:21            return {"error": "No input provided"}22 23        if isinstance(inputs, str):24            inputs = [inputs]25 26        results = []27        for text in inputs:28            encoded = self.tokenizer(29                text,30                return_tensors="pt",31                truncation=True,32                padding="max_length",33                max_length=4096,34            )35            encoded = {k: v.to(self.device) for k, v in encoded.items()}36 37            with torch.no_grad():38                outputs = self.model(**encoded)39 40            raw_score = outputs.logits.squeeze().item()41            clipped_score = min(max(raw_score, 0.0), 1.0)42 43            results.append({"score": round(clipped_score, 4)})44 45        return results46