|
|
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 |
|
|
|
