# handler.py
from typing import Dict, Any, List
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, BitsAndBytesConfig

class EndpointHandler():
    def __init__(self, path=""):
        """
        Carrega o modelo e o tokenizador usando quantização de 4 bits (Q4).
        """
        # 1. Configura a quantização de 4 bits (NF4 com computação em Float16)
        quantization_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.float16,
            bnb_4bit_use_double_quant=True
        )

        # 2. Carrega o tokenizador e o modelo com a configuração de 4 bits
        self.tokenizer = AutoTokenizer.from_pretrained(path)
        self.model = AutoModelForCausalLM.from_pretrained(
            path, 
            quantization_config=quantization_config,
            device_map="auto" # Distribui automaticamente nas GPUs disponíveis
        )
        
        # 3. Encapsula no pipeline padrão do Transformers
        self.pipeline = pipeline(
            "text-generation", 
            model=self.model, 
            tokenizer=self.tokenizer
        )

    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
        """
        Processa as requisições que chegam na API do Inference Endpoint.
        """
        # O Hugging Face padroniza o envio principal dentro do campo "inputs"
        inputs = data.pop("inputs", data)
        parameters = data.pop("parameters", {}) # Parâmetros extras (ex: temperature, max_new_tokens)

        if not inputs:
            return [{"error": "O campo 'inputs' é obrigatório."}]

        # Executa a inferência usando o pipeline configurado no __init__
        prediction = self.pipeline(inputs, **parameters)

        # Retorna o resultado serializável em JSON
        return prediction