from typing import List, Optional, Mapping, Any
from functools import partial
from langchain.llms.base import LLM
from langchain.callbacks.manager import CallbackManagerForLLMRun
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
from transformers import AutoModel, AutoTokenizer
from peft import get_peft_model, LoraConfig, TaskType
import os
import torch

class ChatGLM3(LLM):

    model_path: str
    max_length: int = 8192
    temperature: float = 0.1
    top_p: float = 0.7
    history: List = []
    streaming: bool = True
    model: object = None
    tokenizer: object = None
    """
    def __init__(self, model_path: str, max_length: int = 8192, temperature: float = 0.1, top_p: float = 0.7, history: List = None, streaming: bool = True):
        self.model_path = model_path
        self.max_length = max_length
        self.temperature = temperature
        self.top_p = top_p
        self.history = [] if history is None else history
        self.streaming = streaming
        self.model = None
        self.tokenizer = None
    """

    @property
    def _llm_type(self) -> str:
        return "chatglm3-6B"

    def _call(
        self,
        prompt: str,
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        add_history: bool = False
    ) -> str:
        if self.model is None or self.tokenizer is None:
            raise RuntimeError("Must call `load_model()` to load model and tokenizer!")

        if self.streaming:
            text_callback = partial(StreamingStdOutCallbackHandler().on_llm_new_token, verbose=True)
            resp = self.generate_resp(prompt, text_callback, add_history=add_history)
        else:
            resp = self.generate_resp(self, prompt, add_history=add_history)
        return resp

    def generate_resp(self, prompt, text_callback=None, add_history=True):
        resp = ""
        index = 0
        if text_callback:
            for i, (resp, _) in enumerate(self.model.stream_chat(
                    self.tokenizer,
                    prompt,
                    self.history,
                    max_length=self.max_length,
                    top_p=self.top_p,
                    temperature=self.temperature
            )):
              if add_history:
                if i == 0:
                    self.history += [[prompt, resp]]
                else:
                    self.history[-1] = [prompt, resp]
                text_callback(resp[index:])
                index = len(resp)
        else:
            resp, _ = self.model.chat(
                self.tokenizer,
                prompt,
                self.history,
                max_length=self.max_length,
                top_p=self.top_p,
                temperature=self.temperature
            )
            if add_history:
                self.history += [[prompt, resp]]
        return resp

    def load_model(self):
        if self.model is not None or self.tokenizer is not None:
            return
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_path, trust_remote_code=True)
        self.model = AutoModel.from_pretrained(self.model_path, trust_remote_code=True).half().cuda().eval()

    def load_model_from_checkpoint(self, checkpoint=None):
        if self.model is not None or self.tokenizer is not None:
            return
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_path, trust_remote_code=True)
        self.model = AutoModel.from_pretrained(self.model_path, trust_remote_code=True).half()
        peft_config = LoraConfig(
                task_type=TaskType.CAUSAL_LM,
                inference_mode=False,
                r=8,
                target_modules=['query_key_value'],
                lora_alpha=32,
                lora_dropout=0.1,
        )
        self.model = get_peft_model(self.model, peft_config).to("cuda")
        if checkpoint=="text_classification":
            model_dir = "./output/checkpoint-3000/"
            peft_path = "{}/chatglm-lora.pt".format(model_dir)
            if os.path.exists(peft_path):
                 self.model.load_state_dict(torch.load(peft_path), strict=False)

langchain.callbacks.streaming_stdout.StreamingStdOutCallbackHandler
[((3052, 3122), 'transformers.AutoTokenizer.from_pretrained', 'AutoTokenizer.from_pretrained', (['self.model_path'], {'trust_remote_code': '(True)'}), '(self.model_path, trust_remote_code=True)\n', (3081, 3122), False, 'from transformers import AutoModel, AutoTokenizer\n'), ((3401, 3471), 'transformers.AutoTokenizer.from_pretrained', 'AutoTokenizer.from_pretrained', (['self.model_path'], {'trust_remote_code': '(True)'}), '(self.model_path, trust_remote_code=True)\n', (3430, 3471), False, 'from transformers import AutoModel, AutoTokenizer\n'), ((3589, 3729), 'peft.LoraConfig', 'LoraConfig', ([], {'task_type': 'TaskType.CAUSAL_LM', 'inference_mode': '(False)', 'r': '(8)', 'target_modules': "['query_key_value']", 'lora_alpha': '(32)', 'lora_dropout': '(0.1)'}), "(task_type=TaskType.CAUSAL_LM, inference_mode=False, r=8,\n    target_modules=['query_key_value'], lora_alpha=32, lora_dropout=0.1)\n", (3599, 3729), False, 'from peft import get_peft_model, LoraConfig, TaskType\n'), ((4081, 4106), 'os.path.exists', 'os.path.exists', (['peft_path'], {}), '(peft_path)\n', (4095, 4106), False, 'import os\n'), ((3493, 3559), 'transformers.AutoModel.from_pretrained', 'AutoModel.from_pretrained', (['self.model_path'], {'trust_remote_code': '(True)'}), '(self.model_path, trust_remote_code=True)\n', (3518, 3559), False, 'from transformers import AutoModel, AutoTokenizer\n'), ((3854, 3893), 'peft.get_peft_model', 'get_peft_model', (['self.model', 'peft_config'], {}), '(self.model, peft_config)\n', (3868, 3893), False, 'from peft import get_peft_model, LoraConfig, TaskType\n'), ((1564, 1596), 'langchain.callbacks.streaming_stdout.StreamingStdOutCallbackHandler', 'StreamingStdOutCallbackHandler', ([], {}), '()\n', (1594, 1596), False, 'from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler\n'), ((4152, 4173), 'torch.load', 'torch.load', (['peft_path'], {}), '(peft_path)\n', (4162, 4173), False, 'import torch\n'), ((3144, 3210), 'transformers.AutoModel.from_pretrained', 'AutoModel.from_pretrained', (['self.model_path'], {'trust_remote_code': '(True)'}), '(self.model_path, trust_remote_code=True)\n', (3169, 3210), False, 'from transformers import AutoModel, AutoTokenizer\n')]