微调 BaiChuan13B 进行命名实体识别任务

微调 BaiChuan13B 进行命名实体识别任务 | 极客日志

传统上，一般把 NLP 的研究领域大致分为自然语言理解（NLU）和自然语言生成（NLG）两种。

NLU 侧重于如何理解文本，包括文本分类、命名实体识别、指代消歧、句法分析、机器阅读理解等；

NLG 则侧重于理解文本后如何生成自然文本，包括自动摘要、机器翻译、问答系统、对话机器人等。

但是以 ChatGPT 为代表的大模型出来后，这些传统的 NLP 的细分研究领域基本可以说都失去了独立研究的价值。

为什么呢？因为大模型可以用统一的范式通通将它们搞定，并且效果非常出众。

在之前的例子中，我们演示了使用 QLoRA 算法来对 BaiChuan-13B 实施微调以处理最简单的文本分类任务。在外卖评论数据集上，微调后测试集 acc 由 0.8925 提升到 0.9015 约提升了 1 个百分点。

在本例中，我们使用几乎相同的流程和方法来微调 BaiChuan-13B 以更好地处理命名实体识别任务。

实验结果显示，在 NER 任务上经过微调，我们的 f1-score 取得了不可忽略的提升（0.4313—>0.8768）。

注：跑完本流程需要至少 32G 的 CPU，需要约 2 个小时的训练时间。

在我们正式开始之前，请允许我用简短的话给没有 NLP 基础知识的小伙伴讲解一下什么是命名实体识别。

命名实体识别 NER 任务是 NLP 的一个常见基础任务，它是 Named Entity Recognization 的简称。简单地说，就是识别一个句子中的各种名称实体，诸如：人名，地名，机构等。

例如对于下面这句话：

小明对小红说:"你听说过安利吗？"

其命名实体可以抽取表示如下：

{"人名": ["小明","小红"], "组织": ["安利"]}

〇，预训练模型

我们需要从 Hugging Face 下载 baichuan-13b-chat 的模型。国内可能速度会比较慢，总共有 25G 左右，网速不太好的话，大概可能需要两到三个小时。如果网络不稳定，也可以手动下载全部文件放置到一个文件夹中以便读取。

首先导入必要的库并配置量化参数以节约显存：

import warnings
warnings.filterwarnings('ignore')

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig, AutoModel, BitsAndBytesConfig
from transformers.generation.utils import GenerationConfig
import torch.nn as nn

# 使用 QLoRA 引入的 NF4 量化数据类型以节约显存
model_name_or_path = '../baichuan-13b' # 远程 'baichuan-inc/Baichuan-13B-Chat'

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    llm_int8_threshold=6.0,
    llm_int8_has_fp16_weight=False,
)

tokenizer = AutoTokenizer.from_pretrained(
   model_name_or_path, trust_remote_code=True)

model = AutoModelForCausalLM.from_pretrained(model_name_or_path,
                quantization_config=bnb_config,
                trust_remote_code=True) 

model.generation_config = GenerationConfig.from_pretrained(model_name_or_path)

测试模型的基础对话能力：

from IPython.display import clear_output 
messages = []
messages.append({"role": "user",
                 "content": "世界上第二高的山峰是哪座?"})
response = model.chat(tokenizer,messages=messages,stream=True)
for res in response:
    print(res)
    clear_output(wait=True)

下面我们设计一个 7-shot-prompt 方法，测试一下 BaiChuan13b 的实体抽取能力。

prefix = '''命名实体识别：抽取文本中的 人名，地点，组织 这三类命名实体，并按照 json 格式返回结果。

下面是一些范例：

小明对小红说:"你听说过安利吗？" -> {"人名": ["小明","小红"], "组织": ["安利"]}
现在，每年有几十万中国人到美国访问，几千名中国留学生到美国就学。 -> {"地点": ["中国", "美国"]}
中国是联合国安理会常任理事国之一。 -> {"地点": ["中国"], "组织": ["联合国"]}

请对下述文本进行实体抽取，返回 json 格式。
'''

def get_prompt(text):
    return prefix+text+' -> '

def get_message(prompt,response):
    return [{"role": "user", "content": f'{prompt} -> '},
            {"role": "assistant", "content": response}]

运行测试：

messages  = [{"role": "user", "content": get_prompt("一些摩洛哥球迷已按捺不住，在看台上欢呼雀跃")}]
response = model.chat(tokenizer, messages)
print(response)

输出结果为：

{"地点":["摩洛哥"], "组织":[]}

为了增强 Few-Shot 效果，我们可以将之前的问答记录加入上下文：

messages = messages+[{"role": "assistant", "content": "{'地点': ['摩洛哥']}"}]
messages.extend(get_message("这次轮到北京国安队，不知会不会再步后尘？","{'组织': ['北京国安队']}"))
messages.extend(get_message("革命党人孙中山在澳门成立同盟会分会","{'人名': ['孙中山'], '地名': ['澳门'], '组织': ['同盟会']}"))
messages.extend(get_message("我曾在安徽芜湖市和上海浦东打工。","{'地点': ['安徽芜湖市', '上海浦东']}"))

定义预测函数：

def predict(text,temperature=0.01):
    model.generation_config.temperature=temperature
    response = model.chat(tokenizer, 
                          messages = messages+[{'role':'user','content':f'{text} -> '}])
    return response

predict('杜甫是李白的粉丝。')

输出：

"{'人名': ['杜甫', '李白']}"

我们拿一个开源的中文 NER 数据集来测试一下未经微调，仅仅使用 7-shot-prompt 的预训练模型的效果。

from sklearn.model_selection import train_test_split
import pandas as pd 

df = pd.read_pickle('dfner_13k.pkl')
dfdata,dftest = train_test_split(df,test_size=300,random_state=42)
dftrain,dfval = train_test_split(dfdata,test_size=200,random_state=42)

进行预测循环：

preds = ['' for x in dftest['target']]
for i in tqdm(range(len(preds))):
    preds[i] = predict(dftest['text'].iloc[i])

计算评估指标：

def toset(s):
    try:
        dic = eval(str(s))
        res = []
        for k,v in dic.items():
            for x in v:
                if x:
                    res.append((k,x))
        return set(res)
    except Exception as err:
        print(err)
        return set()

dftest['pred'] = [toset(x) for x in preds]
dftest['gt'] = [toset(x) for x in dftest['target']]
dftest['tp_cnt'] = [len(pred > gt) for pred,gt in zip(dftest['pred'],dftest['gt'])]
dftest['pred_cnt'] = [len(x) for x in dftest['pred']]
dftest['gt_cnt'] = [len(x) for x in dftest['gt']]

precision = sum(dftest['tp_cnt'])/sum(dftest['pred_cnt'])
print('precision = '+str(precision))

recall = sum(dftest['tp_cnt'])/sum(dftest[])
(+(recall))

f1 = *precision*recall/(precision+recall)
(+(f1))

输出结果：

precision = 0.4316109422492401
recall = 0.45151033386327505
f1_score = 0.44133644133644134

微调前 f1_score 为 0.44.

一，准备数据

我们仿照百川模型的 model._build_chat_input 方法来进行 token 编码，同时把需要学习的内容添加 label。

1，token 编码

代码逻辑是将 messages 编码成 token，同时返回 labels。注意 baichuan-13b 通过插入 tokenizer.user_token_id 和 tokenizer.assistant_token_id 来区分用户和机器人会话内容。

import torch 

# 将 messages 编码成 token, 同时返回 labels
# 注意 baichuan-13b 通过插入 tokenizer.user_token_id 和 tokenizer.assistant_token_id 来区分用户和机器人会话内容
# reference@ model._build_chat_input?
def build_chat_input(messages, model=model,
                     tokenizer=tokenizer, 
                     max_new_tokens: int=0):
    max_new_tokens = max_new_tokens or model.generation_config.max_new_tokens
    max_input_tokens = model.config.model_max_length - max_new_tokens
    max_input_tokens = max(model.config.model_max_length // 2, max_input_tokens)
    
    total_input, round_input, total_label, round_label = [], [], [], []
    
    for i, message in enumerate(messages[::-1]):
        content_tokens = tokenizer.encode(message['content'])
        if message['role'] == 'user':
            round_input = [model.generation_config.user_token_id] + content_tokens + round_input
            round_label = [-100]+[-100 for _ in content_tokens]+ round_label
            
            if total_input and len(total_input) + len(round_input) > max_input_tokens:
                break
            else:
                total_input = round_input + total_input
                total_label = round_label + total_label
                if len(total_input) >= max_input_tokens:
                    break
                else:
                    round_input = []
                    round_label = []
                    
        elif message['role'] == 'assistant':
            round_input = [
                model.generation_config.assistant_token_id
            ] + content_tokens + [
                model.generation_config.eos_token_id
            ] + round_input
            
             i==: 
                round_label = [
                    -
                ] + content_tokens + [
                    model.generation_config.eos_token_id
                ]+ round_label
            :
                round_label = [
                    -
                ] + [-  _  content_tokens] + [
                    -
                ]+ round_label
                
        :
             ValueError()
            
    total_input = total_input[-max_input_tokens:]  
    total_label = total_label[-max_input_tokens:]
    
    total_input.append(model.generation_config.assistant_token_id)
    total_label.append(-)
    
     total_input,total_label

2，做数据集

from torch.utils.data import Dataset,DataLoader 
from copy import deepcopy
class MyDataset(Dataset):
    def __init__(self,df,
                 messages
                ):
        self.df = df 
        self.messages = messages
        
    def __len__(self):
        return len(self.df)
        
    def get_samples(self,index):
        samples = []
        d = dict(self.df.iloc[index])
        samples.append(d)
        return samples
    
    def get_messages(self,index):
        samples = self.get_samples(index)
        messages = deepcopy(self.messages)
        for i,d in enumerate(samples):

            messages.append({'role':'user','content':d['text']+' -> '})
            messages.append({'role':'assistant','content':str(d['target'])})
        return messages
        
    def __getitem__(self,index):
        messages = self.get_messages(index)
        input_ids, labels = build_chat_input(messages)
         {:input_ids,:labels}

     ():
        samples = .get_samples(index)
        (samples)

创建训练集和验证集：

ds_train = MyDataset(dftrain,messages)
ds_val = MyDataset(dfval,messages)

3，创建管道

def data_collator(examples: list):
    len_ids = [len(example["input_ids"]) for example in examples]
    longest = max(len_ids) # 之后按照 batch 中最长的 input_ids 进行 padding
    
    input_ids = []
    labels_list = []
    
    for length, example in sorted(zip(len_ids, examples), key=lambda x: -x[0]):
        ids = example["input_ids"]
        labs = example["labels"]
        
        ids = ids + [tokenizer.pad_token_id] * (longest - length)
        labs = labs + [-100] * (longest - length)
        
        input_ids.append(torch.LongTensor(ids))
        labels_list.append(torch.LongTensor(labs))
          
    input_ids = torch.stack(input_ids)
    labels = torch.stack(labels_list)
    return {
        "input_ids": input_ids,
        "labels": labels,
    }

加载 DataLoader：

import torch 
dl_train = torch.utils.data.DataLoader(ds_train,num_workers=2,batch_size=1,
                                       pin_memory=True,shuffle=True,
                                       collate_fn = data_collator)

dl_val = torch.utils.data.DataLoader(ds_val,num_workers=2,batch_size=1,
                                    pin_memory=True,shuffle=False,
                                     collate_fn = data_collator)

for batch in dl_train:
    break 

# 试跑一个 batch
out = model(**batch)
out.loss 

# 采样 300 个 batch 作为一个 epoch，便于较快验证
dl_train.size = 300

二，定义模型

下面我们将使用 QLoRA(实际上用的是量化的 AdaLora）算法来微调 Baichuan-13b 模型。

from peft import get_peft_config, get_peft_model, TaskType
model.supports_gradient_checkpointing = True  
model.gradient_checkpointing_enable()
model.enable_input_require_grads()

model.config.use_cache = False  # silence the warnings. Please re-enable for inference!

查找全连接层以添加 Adapter：

import bitsandbytes as bnb 
def find_all_linear_names(model):
    """
    找出所有全连接层，为所有全连接添加 adapter
    """
    cls = bnb.nn.Linear4bit
    lora_module_names = set()
    for name, module in model.named_modules():
        if isinstance(module, cls):
            names = name.split('.')
            lora_module_names.add(names[0] if len(names) == 1 else names[-1])

    if 'lm_head' in lora_module_names:  # needed for 16-bit
        lora_module_names.remove('lm_head')
    return list(lora_module_names)

from peft import prepare_model_for_kbit_training 
model = prepare_model_for_kbit_training(model)

lora_modules = find_all_linear_names(model)
print(lora_modules)

输出：

['down_proj', 'gate_proj', 'W_pack', 'o_proj', 'up_proj']

配置 AdaLora：

from peft import AdaLoraConfig
peft_config = AdaLoraConfig(
    task_type=TaskType.CAUSAL_LM, inference_mode=False,
    r=16,
    lora_alpha=16, lora_dropout=0.05,
    target_modules= lora_modules
)

peft_model = get_peft_model(model, peft_config)

peft_model.is_parallelizable = True
peft_model.model_parallel = True
peft_model.print_trainable_parameters()

输出：

trainable params: 41,843,040 || all params: 7,002,181,160 || trainable%: 0.5975715144165165

三，训练模型

这里使用自定义的 KerasModel 包装器配合 Accelerator 进行分布式训练支持。

from torchkeras import KerasModel 
from accelerate import Accelerator 

class StepRunner:
    def __init__(self, net, loss_fn, accelerator=None, stage = "train", metrics_dict = None, 
                 optimizer = None, lr_scheduler = None
                 ):
        self.net,self.loss_fn,self.metrics_dict,self.stage = net,loss_fn,metrics_dict,stage
        self.optimizer,self.lr_scheduler = optimizer,lr_scheduler
        self.accelerator = accelerator if accelerator is not None else Accelerator() 
        if self.stage=='train':
            self.net.train() 
        else:
            self.net.eval()
    
    def __call__(self, batch):
        
        #loss
        with self.accelerator.autocast():
            loss = self.net.forward(**batch)[0]

        #backward()
        if self.optimizer is not None and self.stage==:
            .accelerator.backward(loss)
             .accelerator.sync_gradients:
                .accelerator.clip_grad_norm_(.net.parameters(), )
            .optimizer.step()
             .lr_scheduler   :
                .lr_scheduler.step()
            .optimizer.zero_grad()
            
        all_loss = .accelerator.gather(loss).()
        
        
        step_losses = {.stage+:all_loss.item()}
        
        
        step_metrics = {}
        
         .stage==:
             .optimizer   :
                step_metrics[] = .optimizer.state_dict()[][][]
            :
                step_metrics[] = 
         step_losses,step_metrics
    
KerasModel.StepRunner = StepRunner 


 ():
    unwrap_net = accelerator.unwrap_model(.net)
    unwrap_net.save_pretrained(ckpt_path)
    
 ():
     os
    .net.load_state_dict(
        torch.load(os.path.join(ckpt_path,)),strict =)
    .from_scratch = 
    
KerasModel.save_ckpt = save_ckpt 
KerasModel.load_ckpt = load_ckpt

初始化优化器与模型：

optimizer = bnb.optim.adamw.AdamW(peft_model.parameters(),
                                  lr=6e-05,is_paged=True)  #'paged_adamw'
keras_model = KerasModel(peft_model,loss_fn =None,
        optimizer=optimizer) 
ckpt_path = 'baichuan13b_ner'


# keras_model.load_ckpt(ckpt_path) # 支持加载微调后的权重继续训练 (断点续训)
keras_model.fit(train_data = dl_train,
                val_data = dl_val,
                epochs=100,patience=10,
                monitor='val_loss',mode='min',
                ckpt_path = ckpt_path
               )

四，保存模型

为减少 GPU 压力，此处可重启 kernel 释放显存。

import warnings 
warnings.filterwarnings('ignore')

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM,AutoConfig, AutoModel, BitsAndBytesConfig
from transformers.generation.utils import GenerationConfig
import torch.nn as nn
model_name_or_path ='../baichuan-13b'
ckpt_path = 'baichuan13b_ner'
tokenizer = AutoTokenizer.from_pretrained(
    model_name_or_path,
    trust_remote_code=True
)
model_old = AutoModelForCausalLM.from_pretrained(
    model_name_or_path,
    trust_remote_code=True,
    low_cpu_mem_usage=True,
    torch_dtype=torch.float16,
    device_map='auto'
)

合并 LoRA 权重：

from peft import PeftModel

# 可能需要 5 分钟左右
peft_model = PeftModel.from_pretrained(model_old, ckpt_path)
model_new = peft_model.merge_and_unload()

from transformers.generation.utils import GenerationConfig
model_new.generation_config = GenerationConfig.from_pretrained(model_name_or_path)

测试合并后的模型：

from IPython.display import clear_output
messages = []
messages.append({"role": "user",
                 "content": "世界上第二高的山峰是什么？"})
response = model_new.chat(tokenizer,messages=messages,stream=True)
for res in response:
    print(res)
    clear_output(wait=True)

输出：乔戈里峰。世界第二高峰———乔戈里峰西方登山者称其为 k2 峰，海拔高度是 8611 米，位于喀喇昆仑山脉的中巴边境上.

保存最终模型：

save_path = 'baichuan-13b-ner'

tokenizer.save_pretrained(save_path)
model_new.save_pretrained(save_path)

!cp ../baichuan-13b/*.py  baichuan-13b-ner

五，使用模型

为减少 GPU 压力，此处可再次重启 kernel 释放显存。

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM,AutoConfig, BitsAndBytesConfig
from transformers.generation.utils import GenerationConfig
import torch.nn as nn

import warnings
warnings.filterwarnings('ignore')

model_name_or_path = 'baichuan-13b-ner'

... # 加载代码同上
...

我们测试一下微调后的效果。

import pandas as pd 
import numpy as np 
import datasets 
from tqdm import tqdm 

from sklearn.model_selection import train_test_split
import pandas as pd 

df = pd.read_pickle('dfner_13k.pkl')
dfdata,dftest = train_test_split(df,test_size=300,random_state=42)
dftrain,dfval = train_test_split(dfdata,test_size=200,random_state=42)
... # 数据处理同上
...

precision = sum(dftest['tp_cnt'])/sum(dftest['pred_cnt'])
print('precision = '+str(precision))

recall = sum(dftest['tp_cnt'])/sum(dftest['gt_cnt'])
print('recall = '+str(recall))

f1 = 2*precision*recall/(precision+recall)
print('f1_score = '+str(f1))

输出：

precision = 0.9139280125195618
recall = 0.8427128427128427
f1_score = 0.876876876876877

微调后的 f1_score 为 0.8768，相比微调前的 f1_score=0.44，取得了不可忽视的巨大提升。

六，总结与优化建议

本次实验成功展示了利用参数高效微调技术（PEFT）在资源受限环境下适配大模型的能力。通过 QLoRA 结合 AdaLora 策略，我们在保持模型整体性能的同时，仅训练了不到 0.6% 的参数，显著降低了显存占用和训练成本。

在实际应用中，针对命名实体识别任务，还可以进一步优化以下方面：

Prompt Engineering：调整 Few-Shot 示例的数量和质量，确保示例覆盖更多样的实体类型和语境。
数据增强：利用回译或同义词替换增加训练数据的多样性，提高模型泛化能力。
推理加速：在部署阶段，可以使用 vLLM 等推理框架进一步降低延迟，特别是在高并发场景下。
领域适应：如果业务场景涉及特定专业术语，建议在通用语料基础上增加垂直领域的预训练或微调数据。

总之，微调大模型进行特定任务不仅可行，而且随着工具链的成熟，正逐渐成为企业落地 AI 应用的标准路径。开发者应关注模型压缩、量化及推理优化技术，以实现性价比最优的解决方案。

微调 BaiChuan13B 进行命名实体识别任务

〇，预训练模型

一，准备数据

1，token 编码

2，做数据集

3，创建管道

二，定义模型

三，训练模型

四，保存模型

五，使用模型

六，总结与优化建议

更多推荐文章

相关免费在线工具

微调 BaiChuan13B 进行命名实体识别任务

〇，预训练模型

一，准备数据

1，token 编码

2，做数据集

3，创建管道

二，定义模型

三，训练模型

四，保存模型

五，使用模型

六，总结与优化建议

微信扫一扫，关注极客日志

更多推荐文章

相关免费在线工具