yaya36095's picture
Rename handler.py to inference.py
b0e99ab verified
raw
history blame
627 Bytes
from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline
class EndpointHandler:
def __init__(self, path=""):
self.model = AutoModelForSequenceClassification.from_pretrained(
path,
ignore_mismatched_sizes=True # ضروري لتجاوز المشكلة
)
self.tokenizer = AutoTokenizer.from_pretrained(path)
self.pipeline = pipeline("text-classification", model=self.model, tokenizer=self.tokenizer)
def __call__(self, data):
inputs = data.get("inputs") if isinstance(data, dict) else data
return self.pipeline(inputs)