|
|
| 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.
|
| """
|
|
|
| inputs = data.get("inputs", data)
|
|
|
| return self.classifier(inputs)
|
|
|