import torch

from typing import Any, Dict
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig


class EndpointHandler:
    def __init__(self, path=""):
        with torch.autocast('cuda'):
            # load model and tokenizer from path
            self.tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b", padding_side="left")
    
            config = AutoConfig.from_pretrained(path, trust_remote_code=True)
            # config.attn_config['attn_impl'] = 'triton'
            config.init_device = 'cuda:0' # For fast initialization directly on GPU!
            config.max_seq_len = 4096 # (input + output) tokens can now be up to 4096
            
            self.model = AutoModelForCausalLM.from_pretrained(
                path, 
                config,
                torch_dtype=torch.float16,
                trust_remote_code=True
            )
            # self.device = "cuda" if torch.cuda.is_available() else "cpu"
            self.device = 'cuda'

    def __call__(self, data: Dict[str, Any]) -> Dict[str, str]:
        # process input
        inputs = data.pop("inputs", data)
        parameters = data.pop("parameters", None)

        with torch.autocast('cuda'):
            # preprocess
            inputs = self.tokenizer(inputs, return_tensors="pt").to(self.device)
    
            # pass inputs with all kwargs in data
            if parameters is not None:
                outputs = self.model.generate(**inputs, **parameters)
            else:
                outputs = self.model.generate(**inputs)
    
            # postprocess the prediction
            prediction = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
    
            return [{"generated_text": prediction}]