from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
import torch

class EndpointHandler:
    def __init__(self):
        model_id = "klusai/tf2-12b"
        quant = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.bfloat16,
            bnb_4bit_use_double_quant=True,
        )
        self.tokenizer = AutoTokenizer.from_pretrained(model_id)
        self.model = AutoModelForCausalLM.from_pretrained(
            model_id,
            device_map="auto",
            quantization_config=quant,
            torch_dtype=torch.bfloat16,
        )

    def __call__(self, payload: dict):
        prompt = payload.get("inputs") or payload.get("prompt", "")
        max_new = payload.get("max_new_tokens", 1024)
        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
        outputs = self.model.generate(**inputs, max_new_tokens=max_new)
        return {"generated_text": self.tokenizer.decode(outputs[0], skip_special_tokens=True)}