Canthoughts/distilber-emotion-api

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