# handler.py
import json
import os
from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline

class EndpointHandler:
    def __init__(self, path=""):
        # Load the fine-tuned model from the specified path
        self.tokenizer = AutoTokenizer.from_pretrained(path)
        self.model = AutoModelForSequenceClassification.from_pretrained(path)

        # Create a pipeline for text classification
        self.pipeline = pipeline(
            "text-classification",
            model=self.model,
            tokenizer=self.tokenizer
        )

    def __call__(self, data):
        # The data parameter is a dictionary with a 'inputs' key containing the text(s)
        inputs = data.get("inputs", data)
        if isinstance(inputs, str):
            inputs = [inputs]  # Wrap single string in a list

        # Perform inference
        predictions = self.pipeline(inputs)

        # The output should be a list of dictionaries, one for each input
        return predictions