from typing import Any, Dict, List
import torch
import transformers
from transformers import AutoModelForCausalLM, AutoTokenizer

dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] ==8 else torch.float16

class EndpointHandler:
    def __init__(self, path=""):
        self.tokenizer = AutoTokenizer.from_pretrained(path)
        self.model = AutoModelForCausalLM.from_pretrained(path, trust_remote_code=True, revision="main")
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model = self.model.to(self.device)


    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
        prompt = data["inputs"]
        if "config" in data:
          config = data.pop("config", None)
        else:
          config = {'max_new_tokens':100}
        input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids.to(self.device)
        generated_ids = self.model.generate(input_ids, **config)
        generated_text = self.tokenizer.decode(generated_ids[0], skip_special_tokens=True)
        return [{"generated_text": generated_text}]