import os
import torch
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.outputs import OutputKeys
from modelscope.pipelines.nlp.text_generation_pipeline import TextGenerationPipeline
from modelscope.models.base import Model, TorchModel
from modelscope.utils.logger import get_logger
from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
from transformers.generation.utils import GenerationConfig

import torch

from modelscope.models.base import TorchModel
from modelscope.preprocessors.base import Preprocessor
from modelscope.pipelines.base import Model, Pipeline
from modelscope.utils.config import Config
from modelscope.pipelines.builder import PIPELINES
from modelscope.preprocessors.builder import PREPROCESSORS
from modelscope.models.builder import MODELS
# m_path="你的路径/cvx-coder"
# model = AutoModelForCausalLM.from_pretrained(
#     m_path, 
#     device_map="cuda", 
#     torch_dtype="auto", 
#     trust_remote_code=True, 
# )
# tokenizer = AutoTokenizer.from_pretrained(m_path)
# pipe = pipeline(
#     "text-generation", 
#     model=model,
#     tokenizer=tokenizer,
# )
# generation_args = {
#     "max_new_tokens": 2000,
#     "return_full_text": False,
#     "temperature": 0,
#     "do_sample": False,
# }
# content='''my problem is not convex, can i use cvx? if not, what should i do, be specific.'''
# messages = [
#     {"role": "user", "content": content},
# ]
# output = pipe(messages, **generation_args)
# print(output[0]['generated_text'])

@PIPELINES.register_module('text-generation', module_name='cvx-coder-pipe')
class Baichuan7BTextGenerationPipeline(TextGenerationPipeline):
    def __init__(
            self,
            model: Union[Model, str],
            *args,
            **kwargs):
        self.model = Baichuan7BTextGeneration(model) if isinstance(model, str) else model
        super().__init__(model=model, **kwargs)
    
    def preprocess(self, inputs, **preprocess_params) -> Dict[str, Any]:
        return inputs
    
    def _sanitize_parameters(self, **pipeline_parameters):
        return {},pipeline_parameters,{}
    
    # define the forward pass
    def forward(self, inputs: str, **forward_params) -> Dict[str, Any]:
        output = {}
        content=inputs
        messages = [
            {"role": "user", "content": content},
        ]
        outputs = self.model.pipeline(messages, **self.model.generation_args)
       
        output['text'] = outputs[0]['generated_text']
        return output
    
    # format the outputs from pipeline
    def postprocess(self, input, **kwargs) -> Dict[str, Any]:
        return input
    

@MODELS.register_module('text-generation', module_name='cvx-coder')
class Baichuan7BTextGeneration(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)
        self.model = AutoModelForCausalLM.from_pretrained(model_dir, device_map="auto", torch_dtype="auto", trust_remote_code=True)
        # self.model = AutoModelForCausalLM.from_pretrained(model_dir, device_map="auto",trust_remote_code=True)
        self.model.generation_config = GenerationConfig.from_pretrained(model_dir)
        self.model = self.model.eval()
        from transformers import pipeline
        self.pipeline=pipeline("text-generation",model=self.model,tokenizer=self.tokenizer,)
        self.generation_args= {"max_new_tokens": 2000, "return_full_text": False,"temperature": 0, "do_sample": False,}
    
    def forward(self, input: str, *args, **kwargs) -> Dict[str, Any]:
        content = input
        messages = [
            {"role": "user", "content": content},
        ]
        outputs = self.model.pipeline(messages, **self.model.generation_args)
       
        response = outputs[0]['generated_text']
        return {OutputKeys.RESPONSE:response, OutputKeys.HISTORY: ""}
    
    def quantize(self, bits: int):
        self.model = self.model.quantize(bits)
        return self
    
    def infer(self, input, **kwargs):
        content = input
        messages = [
            {"role": "user", "content": content},
        ]
        outputs = self.pipeline(messages, **self.model.generation_args)
       
        response = outputs[0]['generated_text']
        return response


# Tips: usr_config_path is the temporary save configuration location， after upload modelscope hub, it is the model_id
# usr_config_path = '/mnt/workspace/cvx-coder3'
usr_config_path = './'

config = Config({
    "framework": 'pytorch',
    "task": 'text-generation',
    "model": {'type': 'cvx-coder'},
    "pipeline": {"type": "cvx-coder-pipe"},
    "allow_remote": True
})
config.dump('./'+ 'configuration.json')

if __name__ == "__main__":
    from modelscope.models import Model
    from modelscope.pipelines import pipeline
    # model = Model.from_pretrained(usr_config_path)
    input = "Hello, ModelScope!"
    inference = pipeline('text-generation', model=usr_config_path)
    output = inference(input)
    print(output)