from typing import Any, Dict
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import PeftConfig, PeftModel

class EndpointHandler:
  def __init__(self,path=""):
    dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] == 8 else torch.float16

    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_use_double_quant=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=dtype
    )
    config = PeftConfig.from_pretrained(path)
    model = AutoModelForCausalLM.from_pretrained(
                                              config.base_model_name_or_path,
                                              return_dict=True,
                                              quantization_config=bnb_config,
                                              device_map="auto" , 
                                              torch_dtype=dtype,
                                              trust_remote_code=True,
                                              )
    tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
    tokenizer.pad_token = tokenizer.eos_token
    self.tokenizer=tokenizer

    self.model = PeftModel.from_pretrained(model, path)


  def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
      prompt = data.pop("inputs", data)
      prompt = f"""
      ### Instruction:
      Summarize the following conversation.
      ​
      ### Input:
      {prompt}
      ​
      ### Summary:
      """.strip()
      input_ids = self.tokenizer(prompt, return_tensors='pt',truncation=True).input_ids.cuda()
      outputs = self.model.generate(input_ids=input_ids, max_new_tokens=200, )
      output = self.tokenizer.batch_decode(outputs.detach().cpu().numpy(), skip_special_tokens=True)[0][len(prompt):]
      return output