# import os; os.environ["PYTORCH_CUDA_ALLOC_CONF"] = 'max_split_size_mb:1500' # A10
import os; os.environ["PYTORCH_CUDA_ALLOC_CONF"] = 'max_split_size_mb:2000' # A100


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

# class EndpointHandler:
#     def __init__(self, path: str = ""):

#         self.tokenizer = AutoTokenizer.from_pretrained(path, padding_side = "left")
#         self.model = AutoModelForCausalLM.from_pretrained(path, device_map = "auto", torch_dtype=torch.float16)

#     def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
#         """
#         Args:
#             data (:obj:):
#                 includes the input data and the parameters for the inference.
#         Return:
#             A :obj:`list`:. The list contains the answer and scores of the inference inputs
#         """

#         # process input
#         inputs_dict = data.pop("inputs", data)
#         parameters = data.pop("parameters", {})
        
#         prompts = [f"<human>: {prompt}\n<bot>:" for prompt in inputs_dict]

#         self.tokenizer.pad_token = self.tokenizer.eos_token

#         inputs = self.tokenizer(prompts, truncation=True, max_length=2048-512,
#                                 return_tensors='pt', padding=True).to(self.model.device)
#         input_length = inputs.input_ids.shape[1]

#         if parameters.get("deterministic", False):
#             torch.manual_seed(42)

#         outputs = self.model.generate(
#             **inputs, max_new_tokens=512, do_sample=True, temperature=0.7, top_p=0.7, top_k=50
#          )
        
#         output_strs = self.tokenizer.batch_decode(outputs[:, input_length:], skip_special_tokens=True)
        
#         return {"generated_text": output_strs}

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

# class EndpointHandler():
#     def __init__(self, path: str = ""):

#         self.tokenizer = AutoTokenizer.from_pretrained(path, padding_side = "left")
#         self.model = AutoModelForCausalLM.from_pretrained(path, device_map = "auto", torch_dtype=torch.float16)

#     def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
#         """
#         Args:
#             data (:obj:):
#                 includes the input data and the parameters for the inference.
#         Return:
#             A :obj:`list`:. The list contains the answer and scores of the inference inputs
#         """

#         # process input
#         inputs_list = data.pop("inputs", data)
#         parameters = data.pop("parameters", {})
        
#         prompts = [f"<human>: {prompt}\n<bot>:" for prompt in inputs_list]

#         self.tokenizer.pad_token = self.tokenizer.eos_token

#         inputs = self.tokenizer(prompts, truncation=True, max_length=2048-512,
#                                 return_tensors='pt', padding=True).to(self.model.device)
#         input_length = inputs.input_ids.shape[1]

#         if parameters.get("deterministic", False):
#             torch.manual_seed(42)

#         outputs = self.model.generate(
#             **inputs, max_new_tokens=512, do_sample=True, temperature=0.7, top_p=0.7, top_k=50
#          )
        
#         output_strs = self.tokenizer.batch_decode(outputs[:, input_length:], skip_special_tokens=True)
        
#         return {"generated_text": output_strs}

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

class StopWordsCriteria(StoppingCriteria):
    def __init__(self, stop_words, tokenizer):
        self.tokenizer = tokenizer
        self.stop_words = stop_words
        self._cache_str = ''

    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
        self._cache_str += self.tokenizer.decode(input_ids[0, -1])
        for stop_words in self.stop_words:
            if stop_words in self._cache_str:
                return True
        return False


class EndpointHandler():
    def __init__(self, path: str = ""):

        self.tokenizer = AutoTokenizer.from_pretrained(path, padding_side = "left")
        self.tokenizer.pad_token = self.tokenizer.eos_token
        self.model = AutoModelForCausalLM.from_pretrained(path, device_map = "auto", torch_dtype=torch.float16)

    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
        """
        Args:
            data (:obj:):
                includes the input data and the parameters for the inference.
        Return:
            A :obj:`list`:. The list contains the answer and scores of the inference inputs
        """

        # process input
        inputs_list = data.pop("inputs", data)
        parameters = data.pop("parameters", {})
        
        prompts = [f"<human>: {prompt}\n<bot>:" for prompt in inputs_list]

        if parameters.get("EXEC", False):
            exec(parameters['EXEC'])
            del parameters['EXEC']

        if parameters.get("preset_truncation_token"):
          preset_truncation_token_value = parameters["preset_truncation_token"]
          DELIMETER = " "
          prompts = [DELIMETER.join(prompt.split(DELIMETER)[:preset_truncation_token_value]) for prompt in prompts]
          del parameters["preset_truncation_token"]

        with torch.no_grad():
          inputs = self.tokenizer(prompts, truncation=True, max_length=2048-512,
                                  return_tensors='pt', padding=True).to(self.model.device)
          input_length = inputs.input_ids.shape[1]

          if parameters.get("deterministic_seed", False):
              torch.manual_seed(parameters["deterministic_seed"])
              del parameters["deterministic_seed"]
          
          outputs = self.model.generate(
              **inputs, **parameters,
              stopping_criteria=StoppingCriteriaList(
                  [StopWordsCriteria(['\n<human>:'], self.tokenizer)]
                  )
          )
        
        output_strs = self.tokenizer.batch_decode(outputs.sequences[:, input_length:], skip_special_tokens=True)
        output_strs = [output_str.replace("\n<human>:", "") for output_str in output_strs]
        
        torch.cuda.empty_cache()
        return {"generated_text": output_strs}