from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

class EndpointHandler:
    def __init__(self, path=""):
        self.tokenizer = AutoTokenizer.from_pretrained(path)
        self.model = AutoModelForCausalLM.from_pretrained(
            path, torch_dtype=torch.float16, device_map="auto"
        )

    def __call__(self, data):
        inputs = data.pop("inputs", data)
        params = data.pop("parameters", {})
        encoded = self.tokenizer(inputs, return_tensors="pt").to(self.model.device)
        output = self.model.generate(
            **encoded,
            max_new_tokens=params.get("max_new_tokens", 600),
            temperature=params.get("temperature", 0.1),
            repetition_penalty=params.get("repetition_penalty", 1.1),
            do_sample=True,
        )
        return self.tokenizer.decode(output[0][encoded["input_ids"].shape[1]:], skip_special_tokens=True)