| | from transformers import AutoTokenizer, AutoModelForSequenceClassification |
| | import torch |
| |
|
| | class EndpointHandler: |
| | def __init__(self, path=""): |
| | |
| | self.tokenizer = AutoTokenizer.from_pretrained(path) |
| | self.model = AutoModelForSequenceClassification.from_pretrained(path) |
| | self.model.eval() |
| | self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| | self.model.to(self.device) |
| |
|
| | def __call__(self, data): |
| | """ |
| | This method is called when the endpoint receives a request. |
| | Expected input: { "inputs": "some string" } or { "inputs": ["a", "b", ...] } |
| | """ |
| | inputs = data.get("inputs", None) |
| |
|
| | if inputs is None: |
| | return {"error": "No input provided"} |
| |
|
| | if isinstance(inputs, str): |
| | inputs = [inputs] |
| |
|
| | results = [] |
| | for text in inputs: |
| | encoded = self.tokenizer( |
| | text, |
| | return_tensors="pt", |
| | truncation=True, |
| | padding="max_length", |
| | max_length=4096, |
| | ) |
| | encoded = {k: v.to(self.device) for k, v in encoded.items()} |
| |
|
| | with torch.no_grad(): |
| | outputs = self.model(**encoded) |
| |
|
| | raw_score = outputs.logits.squeeze().item() |
| | clipped_score = min(max(raw_score, 0.0), 1.0) |
| |
|
| | results.append({"score": round(clipped_score, 4)}) |
| |
|
| | return results |
