File size: 2,634 Bytes
52bae7f d1753e0 52bae7f 6f71ef9 52bae7f d1753e0 52bae7f d1753e0 52bae7f d1753e0 52bae7f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 | from transformers import (
AutoConfig,
AutoTokenizer,
AutoModelForSequenceClassification,
)
import torch
from typing import Dict, Any, List
class EndpointHandler:
def __init__(self, path: str = ""):
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
config = AutoConfig.from_pretrained(path)
config.num_labels = 1
self.tokenizer = AutoTokenizer.from_pretrained(path, use_fast=True)
self.model = AutoModelForSequenceClassification.from_pretrained(
path,
config=config,
ignore_mismatched_sizes=True,
)
self.model.to(self.device)
self.model.eval()
self.threshold = float(getattr(self.model.config, "threshold_F1", 0.5))
self.num_labels = int(getattr(self.model.config, "num_labels", 1))
def _predict_probs(self, inputs_text: List[str]) -> List[float]:
encoded = self.tokenizer(
inputs_text,
return_tensors="pt",
truncation=True,
padding=True,
max_length=512,
)
encoded = {k: v.to(self.device) for k, v in encoded.items()}
with torch.no_grad():
logits = self.model(**encoded).logits
if self.num_labels == 1 or logits.shape[-1] == 1:
logits = logits.squeeze(-1)
probs = torch.sigmoid(logits).detach().cpu().tolist()
else:
probs = torch.softmax(logits, dim=-1)[:, 1].detach().cpu().tolist()
if isinstance(probs, float):
probs = [probs]
return probs
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
inputs_text = data.get("inputs", "")
if isinstance(inputs_text, str):
inputs_text = [inputs_text]
elif not isinstance(inputs_text, list):
raise ValueError("`inputs` must be a string or a list of strings.")
inputs_text = [
str(x).strip()
for x in inputs_text
if x is not None and str(x).strip() != ""
]
if not inputs_text:
return []
probs = self._predict_probs(inputs_text)
return [
{
"text": text,
"label": "CAUSAL" if prob >= self.threshold else "NON-CAUSAL",
"score": round(float(prob), 6),
"p_causal": round(float(prob), 6),
"threshold": round(float(self.threshold), 6),
}
for text, prob in zip(inputs_text, probs)
] |