# Copyright (c) 2022 Zhipu.AI
import torch
from modelscope.pipelines.builder import PIPELINES
from modelscope.pipelines.base import Pipeline
from modelscope.utils.constant import Tasks
from modelscope.models.builder import MODELS
from modelscope.utils.logger import get_logger
from modelscope.models.base import Model, TorchModel
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers import PreTrainedModel, PreTrainedTokenizer
import numpy as np
from typing import Union, Dict, Any

PROMPT_FORMAT = """Below is an instruction that describes a task. Write a response that appropriately completes the request.

### Instruction:
{instruction}

### Response:
"""

def generate_response(instruction: str, *, model: PreTrainedModel, tokenizer: PreTrainedTokenizer, 
                      do_sample: bool = True, max_new_tokens: int = 256, top_p: float = 0.92, top_k: int = 0, device=None, **kwargs) -> str:
    input_ids = tokenizer(PROMPT_FORMAT.format(instruction=instruction), return_tensors="pt").input_ids
    # each of these is encoded to a single token
    response_key_token_id = tokenizer.encode("### Response:")[0]
    end_key_token_id = tokenizer.encode("### End")[0]
    input_ids = input_ids.to(torch.device(device))
    model = model.to(torch.device(device))
    gen_tokens = model.generate(input_ids, pad_token_id=tokenizer.pad_token_id, eos_token_id=end_key_token_id,
                                do_sample=do_sample, max_new_tokens=max_new_tokens, top_p=top_p, top_k=top_k, **kwargs)[0].cpu()

    # find where the response begins
    response_positions = np.where(gen_tokens == response_key_token_id)[0]

    if len(response_positions) >= 0:
        response_pos = response_positions[0]
        
        # find where the response ends
        end_pos = None
        end_positions = np.where(gen_tokens == end_key_token_id)[0]
        if len(end_positions) > 0:
            end_pos = end_positions[0]

        return tokenizer.decode(gen_tokens[response_pos + 1 : end_pos]).strip()

    return None


@PIPELINES.register_module(Tasks.text_generation, module_name='DollyV16b-text-generation')
class DollyV16bTextGenerationPipeline(Pipeline):
    def __init__(self,
                 model: Union[Model, str],
                 *args,
                 **kwargs):
        device = kwargs.get('device')
        model = DollyV16bTextGeneration(model_dir=model,device=device) if isinstance(model, str) else model
        self.model = model
        self.model.eval()
        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='DollyV16b')
class DollyV16bTextGeneration(TorchModel):

    def __init__(self, model_dir: str, device=None, *args, **kwargs):
        super().__init__(model_dir, *args, **kwargs)
        self.logger = get_logger()
        self.device = device
        # loading tokenizer
        self.tokenizer = AutoTokenizer.from_pretrained(model_dir, padding_side="left")
        # loading model
        self.model = AutoModelForCausalLM.from_pretrained(model_dir, trust_remote_code=True)



    
    def forward(self, input: Dict) -> Dict:
        output = {}
        res = generate_response(input,model=self.model,tokenizer=self.tokenizer,device=self.device)
        output['text'] = res
        return output
    
    def quantize(self, bits: int):
        self.model = self.model.quantize(bits)
        return self