# handler.py
from typing import Any, Dict
import torch
from transformers import pipeline

class EndpointHandler:
    def __init__(self, path: str = ""):
        self.pipe = pipeline(
            "text-generation",
            model=path,
            torch_dtype=torch.bfloat16,
            device_map="auto",
        )

    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
        inputs = data.pop("inputs", None)
        if inputs is None:
            return {"error": "Missing `inputs`"}

        params = data.pop("parameters", {})

        # sensible defaults
        params.setdefault("max_new_tokens", 512)
        params.setdefault("temperature", 0.7)
        params.setdefault("do_sample", True)
        params.setdefault("return_full_text", False)

        output = self.pipe(inputs, **params)

        return {
            "generated_text": output[0]["generated_text"]
        }