import os
from typing import Union, Dict, Any
from modelscope.pipelines.builder import PIPELINES
from modelscope.models.builder import MODELS
from modelscope.utils.constant import Tasks
from modelscope.pipelines.base import Pipeline
from modelscope.models.base import Model, TorchModel
from modelscope.utils.logger import get_logger
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers import pipeline


os.environ['CUDA_VISIBLE_DEVICES'] = "0"

@PIPELINES.register_module(Tasks.text_generation, module_name='opt-125-text-generation-pipe')
class Opt125TextGenerationPipeline(Pipeline):
    def __init__(
            self,
            model: Union[Model, str],
            *args,
            **kwargs):
        model = Opt125TextGeneration(model) if isinstance(model, str) else model
        super().__init__(model=model, **kwargs)
    
    def preprocess(self, inputs, **preprocess_params) -> Dict[str, Any]:
        return inputs
    
    # define the forward pass
    def forward(self, inputs: Dict, **forward_params) -> Dict[str, Any]:
        return self.model(inputs)
    
    # format the outputs from pipeline
    def postprocess(self, input, **kwargs) -> Dict[str, Any]:
        return input

@MODELS.register_module(Tasks.text_generation, module_name='opt-125')
class Opt125TextGeneration(TorchModel):
    def __init__(self, model_dir=None, *args, **kwargs):
        super().__init__(model_dir, *args, **kwargs)
        self.logger = get_logger()
        # loading tokenizer
        self.tokenizer = AutoTokenizer.from_pretrained(model_dir,trust_remote_code=True)
        self.model = AutoModelForCausalLM.from_pretrained(model_dir, device_map="auto",trust_remote_code=True)
        self.generator = pipeline('text-generation', model=model_dir)
        self.model = self.model.eval()
    
    def forward(self,input: Dict, *args, **kwargs) -> Dict[str, Any]:
        output = {}
        res = self.generator(input)
        output['text'] = res[0]['generated_text']
        return output
    
    def quantize(self, bits: int):
        self.model = self.model.quantize(bits)
        return self