File size: 833 Bytes
af04c87 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 | # handler.py
from typing import Any, Dict, List
from transformers import pipeline
class EndpointHandler:
def __init__(self, path: str = ""):
"""
Load your model and create a Hugging Face pipeline.
'path' is the local folder or repo name containing your model.
"""
self.classifier = pipeline("text-classification", model=path)
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
"""
Called on each inference request.
Expects a dict with an "inputs" key (string or list of strings).
Returns the pipeline output as a list of dicts.
"""
# Extract inputs; if they passed raw string, handle that too
inputs = data.get("inputs", data)
# Run inference
return self.classifier(inputs)
|