ChatGLM-6B：自定义数据集和模型微调！

 Datawhale干货 

开源：ChatGLM，编辑：Coggle数据科学

ChatGLM-6B介绍

ChatGLM-6B是一个开源的、支持中英双语问答的对话语言模型，基于 General Language Model (GLM)架构，具有 62 亿参数。

具体代码和权重可以从👇获取：

https://github.com/THUDM/ChatGLM-6B/
https://huggingface.co/THUDM/chatglm-6b

ChatGLM-6B拥有不同版本，暂时来看是由权重保存不同格式带来区别，但对模型的效果应该没有影响。

量化等级	最低 GPU 显存（推理）	最低 GPU 显存（高效参数微调）
FP16（无量化）	13 GB	14 GB
INT8	8 GB	9 GB
INT4	6 GB	7 GB

ChatGLM-6B 自定义数据集

https://github.com/THUDM/ChatGLM-6B/tree/main/ptuning

ChatGLM-6B需要将数据集定义句子对的形式。这里的需要考虑下游任务，以及prompt的构造案例。

下面为从关键字到句子的案例：

{
    "content": "类型#上衣*版型#宽松*版型#显瘦*图案#线条*衣样式#衬衫*衣袖型#泡泡袖*衣款式#抽绳",
    "summary": "这件衬衫的款式非常的宽松，利落的线条可以很好的隐藏身材上的小缺点，穿在身上有着很好的显瘦效果。领口装饰了一个可爱的抽绳，漂亮的绳结展现出了十足的个性，配合时尚的泡泡袖型，尽显女性甜美可爱的气息。"
}

如果需要进行多轮对话，则需要考虑到上下对话，汇总在history中：

{
  "prompt": "长城h3风扇不转。继电器好的。保险丝好的传感器新的风扇也新的这是为什么。就是继电器缺一个信号线", 
  "response": "用电脑能读数据流吗？水温多少", 
  "history": []
}

{
  "prompt": "95", 
  "response": 
  "上下水管温差怎么样啊？空气是不是都排干净了呢？", 
  "history": [["长城h3风扇不转。继电器好的。保险丝好的传感器新的风扇也新的这是为什么。就是继电器缺一个信号线", "用电脑能读数据流吗？水温多少"]]
}

ChatGLM-6B 微调过程

https://github.com/THUDM/ChatGLM-6B/tree/main/ptuning

在训练过程设置自定义的数据集路径，然后也需要设置如下的超参数。

pre_seq_len=128
learning_rate=2e-2
quantization_bit=4
per_device_train_batch_size=16
gradient_accumulation_steps=1

PRE_SEQ_LEN和LR分别是 soft prompt 长度和训练的学习率，可以进行调节以取得最佳的效果。

P-Tuning-v2 方法会冻结全部的模型参数，可通过调整 quantization_bit 来被原始模型的量化等级。

ChatGLM-6B 加载并预测

当模型训练完成，接下来就可以加载模型并进行预测。加载的代码如下：

from transformers import AutoConfig, AutoModel, AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True)

config = AutoConfig.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True, pre_seq_len=128)
model = AutoModel.from_pretrained("THUDM/chatglm-6b", config=config, trust_remote_code=True)
# 自定义训练得到模型的路径
prefix_state_dict = torch.load(os.path.join(CHECKPOINT_PATH, "pytorch_model.bin"))
new_prefix_state_dict = {}
for k, v in prefix_state_dict.items():
    if k.startswith("transformer.prefix_encoder."):
        new_prefix_state_dict[k[len("transformer.prefix_encoder."):]] = v

model.transformer.prefix_encoder.load_state_dict(new_prefix_state_dict)

如果需要加载的是旧 Checkpoint（包含ChatGLM-6B以及 PrefixEncoder 参数），或者进行的是全参数微调，则直接加载整个 Checkpoint：

model = AutoModel.from_pretrained(CHECKPOINT_PATH, trust_remote_code=True)

之后可以进行量化，也可以直接使用：

# Comment out the following line if you don't use quantization
model = model.quantize(4)
model = model.half().cuda()
model.transformer.prefix_encoder.float()
model = model.eval()

response, history = model.chat(tokenizer, "你好", history=[])

一起“点赞”三连↓