from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline

class EndpointHandler:
    def __init__(self, model_path: str, task="text-generation"):
        self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
        self.model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True)
        self.pipe = pipeline(task=task, model=self.model, tokenizer=self.tokenizer)

    def __call__(self, inputs: dict) -> dict:
        prompt = inputs.get("inputs", "")
        params = inputs.get("parameters", {})
        outputs = self.pipe(prompt, **params)
        return outputs