#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Project      : AI.  @by PyCharm
# @File         : utils
# @Time         : 2023/4/20 12:50
# @Author       : betterme
# @WeChat       : meutils
# @Software     : PyCharm
# @Description  :
import torch
from transformers import AutoTokenizer, AutoModel, AutoConfig

from meutils.pipe import *


def auto_configure_device_map(num_gpus: int) -> Dict[str, int]:
    # transformer.word_embeddings 占用1层
    # transformer.final_layernorm 和 lm_head 占用1层
    # transformer.layers 占用 28 层
    # 总共30层分配到num_gpus张卡上
    num_trans_layers = 28
    per_gpu_layers = 30 / num_gpus

    # bugfix: 在linux中调用torch.embedding传入的weight,input不在同一device上,导致RuntimeError
    # windows下 model.device 会被设置成 transformer.word_embeddings.device
    # linux下 model.device 会被设置成 lm_head.device
    # 在调用chat或者stream_chat时,input_ids会被放到model.device上
    # 如果transformer.word_embeddings.device和model.device不同,则会导致RuntimeError
    # 因此这里将transformer.word_embeddings,transformer.final_layernorm,lm_head都放到第一张卡上
    device_map = {'transformer.word_embeddings': 0,
                  'transformer.final_layernorm': 0, 'lm_head': 0}

    used = 2
    gpu_target = 0
    for i in range(num_trans_layers):
        if used >= per_gpu_layers:
            gpu_target += 1
            used = 0
        assert gpu_target < num_gpus
        device_map[f'transformer.layers.{i}'] = gpu_target
        used += 1

    return device_map


def llm_load(model_name_or_path="THUDM/chatglm-6b", device='cpu', device_map: Optional[Dict[str, int]] = None,
             **kwargs):
    tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, trust_remote_code=True)
    model = AutoModel.from_pretrained(model_name_or_path, trust_remote_code=True, **kwargs)

    if torch.cuda.is_available() and device.lower().startswith("cuda"):
        # 根据当前设备GPU数量决定是否进行多卡部署
        num_gpus = torch.cuda.device_count()
        if num_gpus < 2 and device_map is None:
            model = model.half().cuda()
        else:
            from accelerate import dispatch_model
            # 可传入device_map自定义每张卡的部署情况
            if device_map is None:
                device_map = auto_configure_device_map(num_gpus)

            model = dispatch_model(model, device_map=device_map).half().cuda()

    else:
        model = model.float().to(device)

    return model.eval(), tokenizer


def llm_load4chat(model_name_or_path="THUDM/chatglm-6b", device='cpu', stream=True, **kwargs):
    model, tokenizer = llm_load(model_name_or_path, device, **kwargs)
    if stream:
        return partial(model.stream_chat, tokenizer=tokenizer)
    else:
        return partial(model.chat, tokenizer=tokenizer)


if __name__ == '__main__':
    a = llm_load("/Users/betterme/PycharmProjects/AI/CHAT_MODEL/chatglm")
