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

class PreTrainedPipeline():
    def __init__(self, path=""):
        path = "oleksandrfluxon/mpt-7b-instruct-2"
        print("===> path", path)
        
        config = transformers.AutoConfig.from_pretrained(path, trust_remote_code=True)
        config.max_seq_len = 4096 # (input + output) tokens can now be up to 4096

        print("===> loading model")
        model = transformers.AutoModelForCausalLM.from_pretrained(
          path,
          config=config,
          torch_dtype=torch.bfloat16, # Load model weights in bfloat16
          trust_remote_code=True,
          load_in_4bit=True, # Load model in the lowest 4-bit precision quantization
        )
        print("===> model loaded")

        tokenizer = transformers.AutoTokenizer.from_pretrained('EleutherAI/gpt-neox-20b', padding_side="left", device_map="auto")
            
        self.pipeline = transformers.pipeline('text-generation', model=model, tokenizer=tokenizer)
        print("===> init finished")

    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
        """
       data args:
            inputs (:obj: `str`)
            parameters (:obj: `str`)
      Return:
            A :obj:`str`: todo
        """
        # get inputs
        inputs = data.pop("inputs",data)
        parameters = data.pop("parameters", {})
        date = data.pop("date", None)
        print("===> inputs", inputs)
        print("===> parameters", parameters)

        result = self.pipeline(inputs, **parameters)
        print("===> result", result)

        return result