前记：

BERT(Bidirectional Encoder Representation from Transformers)是2018年10月Google AI研究院提出的一种预训练-微调大模型，和GPT属于是同一时期的产物。BERT在多项自然语言处理任务上都实现了SOTA，抛弃了传统模型，在一种任务上需要修改模型架构的方式，而是同一个模型架构在不需要修改的情况下，应用于多个任务上。因此，一经问世就引起了学术界极大地关注，甚至超过了GPT。

BERT的网络架构使用的是基于Transformer中的Encoder部分，由于其使用Attention的关系，从而避免了传统网络结构，诸如RNN，LSTM，GRU等的遗忘问题，可以捕获一个长序列不同token的相关性问题。

Bert架构图如下：
Bert架构图

网络结构介绍：

这里重点介绍一下Bert的嵌入编码部分，后面部分就是Encoder模块的堆叠了，并没有什么新奇的。
Bert的嵌入编码有三部分，分别是token Embedding，position Embedding, segment Embedding，之所以这里需要segment Embedding，是因为在Bert的预训练任务中有两个训练任务一个是掩码填充预测，一个NSP(Next Sentence Prediction)预测。

NER任务：

命名实体识别——Named Entity Recognition——简称NER，是指从文本中识别出具有特定意义的实体，主要包括人名、地名、机构名、专有名词等。 简单举例： 识别出的实体结果： 结果显示就把举例文本中的具有特定意义的实体，机构名和人名以及它们的具体位置都识别出来了。 以上就是一个NER任务的实际样例，在做具体的项目的时候，往往涉及到的实体类型更多，有可能成百上千，然后数据量更多。

使用Bert用于NER任务

所使用到的数据集链接：数据集

首先是数据集的处理：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
import torch
import json

class NerDataset(torch.utils.data.Dataset):
    def __init__(self,filename,tokenizer):
        self.tokenizer=tokenizer
        with open(filename,'r',encoding='utf-8') as f:
            self.tokens_arr,self.labels_arr=self._preprocess(json.load(f)) # 读取json文件内容并进行预处理
    def _preprocess(self,data):
        tokens_arr,labels_arr=[],[] # 对标注数据的token 和 token_label_list进行处理
        for item in data:
            text=item['data']['text'] # 原始文本
            annotations=item['annotations'] # 标签数据

            ch_tokens_list=[]
            ch_labels_list=[]
            for ch in text: # tokenizer by char
                ch_tokens=self.tokenizer([ch],add_special_tokens=False)
                assert len(ch_tokens.input_ids[0])<=1
                # print(ch_tokens.input_ids[0])
                ch_tokens_list.append(ch_tokens.input_ids[0])
                ch_labels_list.append(['O']*len(ch_tokens.input_ids[0])) # 全部复制为'o'
            for anno in annotations[0]['result']:
                anno=anno['value']
                start,end,text,label=anno['start'],anno['end'],anno['text'],anno['labels'][0]
                first=True
                for i in range(start,end): # 对数据进行标注
                    for j in range(len(ch_labels_list[i])):
                        if first:
                            ch_labels_list[i][j]=label+'-B' # 开始片段
                            first=False
                        else:
                            ch_labels_list[i][j]=label+'-I' # 中间片段
            item_tokens=[]
            item_labels=[]
            for i in range(len(ch_tokens_list)):
                item_tokens.extend(ch_tokens_list[i]) # 一个token样本
                item_labels.extend(ch_labels_list[i]) # 对应的一个标签样本
            tokens_arr.append(item_tokens)
            labels_arr.append(item_labels)
        return tokens_arr,labels_arr # 返回所有的token和token标签

    def __getitem__(self,idx):
        return self.tokens_arr[idx],self.labels_arr[idx]

    def __len__(self):
        return len(self.tokens_arr)

if __name__=='__main__':
    # with open('/content/drive/MyDrive/Colab Notebooks/ner-bert-crf/data/train_label.json','r',encoding='utf-8') as f:
      # print(type(json.load(f)[0]))
      # print(json.load(f)[0])
    from modelscope import AutoTokenizer
    tokenizer=AutoTokenizer.from_pretrained('google-bert/bert-base-chinese') # 加载分词器
    ds=NerDataset('/content/drive/MyDrive/Colab Notebooks/ner-bert-crf/data/train_label.json',tokenizer)  # 加载标注数据
    tokens,labels=ds[0]
    print(tokens,len(tokens)) 
    print(tokenizer.decode(tokens))
    print(labels,len(labels)) # 标注结果

output:

1
2
3
[1745, 711, 1266, 776, 1920, 2110, 1744, 2157, 1355, 2245, 4777, 4955, 7368, 3136, 2956, 510, 704, 1744, 5307, 3845, 4777, 4955, 704, 2552, 712, 818, 2001, 3817, 1745, 120, 1358, 6393, 5442, 897, 1745] 35
图 为 北 京 大 学 国 家 发 展 研 究 院 教 授 、 中 国 经 济 研 究 中 心 主 任 姚 洋 图 / 受 访 者 供 图
['O', 'O', 'COUNTRY-B', 'COUNTRY-I', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'O', 'PERSON-B', 'PERSON-I', 'O', 'O', 'O', 'O', 'O', 'O', 'O'] 35

模型构建：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
import torch
from torchcrf import CRF
from modelscope import BertModel

class NerBertCrfModel(torch.nn.Module):
    def __init__(self,num_tags):
        super().__init__()
        self.bert=BertModel.from_pretrained('google-bert/bert-base-chinese') # 加载Bert中文模型
        self.hidden=torch.nn.Linear(in_features=768,out_features=num_tags) # 输出设置tag类别数
        self.crf=CRF(num_tags=num_tags,batch_first=True)
        for p in self.bert.parameters(): # 冻结bert原始预训练的参数
            p.requires_grad=False
            
    def forward(self,input_ids,attention_mask,token_type_ids,labels):
        emissions=self.bert(input_ids,attention_mask,token_type_ids)
        # print(emissions.last_hidden_state.shape) # (2,20,768)
        emissions=self.hidden(emissions.last_hidden_state) # (2,20,5)
        return -self.crf(emissions=emissions,tags=labels,mask=attention_mask) # loss 对数似然值取相反数，作为损失函数
    
    def predict(self,input_ids,attention_mask,token_type_ids):
        emissions=self.bert(input_ids,attention_mask,token_type_ids)
        emissions=self.hidden(emissions.last_hidden_state)
        return self.crf.decode(emissions=emissions,mask=attention_mask)

if __name__=='__main__':
    model=NerBertCrfModel(num_tags=5)
    input_ids=torch.randint(0,100,size=(2,20))
    attention_mask=torch.ones_like(input_ids).bool()
    token_type_ids=torch.zeros_like(input_ids)
    labels=torch.randint(0,5,size=(2,20))
    loss=model(input_ids,attention_mask,token_type_ids,labels)
    predict=model.predict(input_ids,attention_mask,token_type_ids)
    print(predict)

output:

1
2
# torch.Size([2, 20, 768])
[[2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2], [2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2]]

训练代码部分：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
# from dataset import NerDataset
# from model import NerBertCrfModel
from modelscope import AutoTokenizer
import torch
import functools
import os
import tqdm

device='cuda' if torch.cuda.is_available() else 'cpu'

labels={'COUNTRY-B':0,'COUNTRY-I':1,'PERSON-B':2,'PERSON-I':3,'O':4} # 标签ID映射

def collate_fn(batch,tokenizer): # 对数据进行批量处理，填充成相同长度
    bos_id=tokenizer.convert_tokens_to_ids(['[CLS]']) # 开始符号ID
    eos_id=tokenizer.convert_tokens_to_ids(['[SEP]']) # 结束符号ID
    pad_id=tokenizer.convert_tokens_to_ids(['[PAD]']) # 填充符号ID
    batch_x,batch_y,batch_attn_mask,batch_token_type_ids=[],[],[],[]
    max_seq_len=0
    for sample in batch: # 训练数据token处理
        x=bos_id+sample[0]+eos_id
        batch_x.append(x)
        if len(x)>max_seq_len:
            max_seq_len=len(x)
        y=[labels['O']]+[labels[label] for label in sample[1]]+[labels['O']]
        batch_y.append(y)
    # padding

    print('max:',max_seq_len)
    for i,x in enumerate(batch_x): # batch_attn_mask是填充掩码，标注出PAD填充掩码
        batch_attn_mask.append([1]*len(x)+[0]*(max_seq_len-len(x)))
        batch_y[i].extend([labels['O']]*(max_seq_len-len(x)))
        x.extend(pad_id*(max_seq_len-len(x))) # 填充 x
        batch_token_type_ids.append([0]*len(x)) # sentence A 分段掩码
    return torch.tensor(batch_x,dtype=torch.long),torch.tensor(batch_attn_mask,dtype=torch.bool),torch.tensor(batch_token_type_ids,dtype=torch.long),torch.tensor(batch_y,dtype=torch.long)

if __name__=='__main__':
    tokenizer=AutoTokenizer.from_pretrained('google-bert/bert-base-chinese') # 加载分词器
    model=NerBertCrfModel(num_tags=len(labels)).to(device) # 加载模型
    try:
        model.load_state_dict(torch.load('model.pt')) # 加载模型权重
    except:
        pass
    optimizer=torch.optim.Adam([p for p in model.parameters() if p.requires_grad],lr=6e-5)

    dataset=NerDataset('/content/drive/MyDrive/Colab Notebooks/ner-bert-crf/data/train_label.json',tokenizer)
    dataloader=torch.utils.data.DataLoader(dataset,batch_size=16,shuffle=True,persistent_workers=True,num_workers=2,collate_fn=functools.partial(collate_fn,tokenizer=tokenizer))

    steps=0
    epoch=0
    model.train()
    while True:
        pbar=tqdm.tqdm(dataloader,total=len(dataloader),ncols=100)
        for batch_x,batch_attn_mask,batch_token_type_ids,batch_y in pbar:
            batch_x,batch_attn_mask,batch_token_type_ids,batch_y=batch_x.to(device),batch_attn_mask.to(device),batch_token_type_ids.to(device),batch_y.to(device)
            loss=model(batch_x,batch_attn_mask,batch_token_type_ids,batch_y)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            steps+=1
            pbar.set_description(f'Epoch {epoch}')
            pbar.set_postfix(loss=loss.item(),steps=steps)

            # checkpoint
            if steps%300==0:
                torch.save(model.state_dict(),'.model.pt')
                os.replace('.model.pt','model.pt')
            if epoch == 100:
              torch.save(model.state_dict(),'.model.pt')
              os.replace('.model.pt','model.pt')
              print('finish...')
              break
        epoch+=1

evolution：这部分主要是对预测的后处理结果，处理后返回后处理的识别分类元组。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
import torch
from torchcrf import CRF
# from model import NerBertCrfModel
from modelscope import AutoTokenizer

device='cuda' if torch.cuda.is_available() else 'cpu'
labels={'COUNTRY-B':0,'COUNTRY-I':1,'PERSON-B':2,'PERSON-I':3,'O':4}
labels_rev=dict((v,k) for k,v in labels.items())

tokenizer=AutoTokenizer.from_pretrained('google-bert/bert-base-chinese')
model=NerBertCrfModel(num_tags=len(labels)).to(device)
model.load_state_dict(torch.load('/content/drive/MyDrive/Colab Notebooks/model.pt'))
model.eval()

def ner(s):
    tokens=[]
    for ch in s:
        ret=tokenizer([ch],add_special_tokens=False)
        tokens=tokens+ret.input_ids[0]

    bos_id=tokenizer.convert_tokens_to_ids(['[CLS]'])
    eos_id=tokenizer.convert_tokens_to_ids(['[SEP]'])

    input_ids=torch.tensor(bos_id+tokens+eos_id,dtype=torch.long).to(device)
    attn_mask=torch.tensor([1]*len(input_ids),dtype=torch.bool).to(device)
    type_ids=torch.tensor([0]*len(input_ids),dtype=torch.long).to(device)

    pred=model.predict(input_ids.unsqueeze(0),attn_mask.unsqueeze(0),type_ids.unsqueeze(0))
    # ignore [CLS] and [SEP]
    input_ids=input_ids[1:-1]
    pred=pred[0][1:-1]
    start=None
    entity=''
    ner_result=[]
    for i in range(len(pred)):
        pred_label=labels_rev[pred[i]]
        pred_label_splits=pred_label.split('-')
        pred_label_first=pred_label_splits[0]
        pred_label_second='' if len(pred_label_splits)<=1 else pred_label_splits[1]
        if start is None and pred_label_second=='B': # entiry start
            start=i
            entity=pred_label_first
        elif start is not None and (pred_label_first!=entity or (pred_label_first==entity and pred_label_second=='B')):
            entity_value=[tokenizer.convert_ids_to_tokens([id])[0] for id in input_ids[start:i]]
            ner_result.append((start,i-1,entity,''.join(entity_value)))
            if pred_label_second=='B':
                start=i
                entity=pred_label_first
            else:
                start=None
                entity=''
    if start is not None:
        entity_value=[tokenizer.convert_ids_to_tokens([id])[0] for id in input_ids[start:]]
        ner_result.append((start,len( )-1,entity,''.join(entity_value)))
    return ner_result

s='姚洋：我方的这种做法是对的。特朗普发动关税战以来，我就认为我们要警惕落入特朗普设定的圈圈，或者说是某种意义上的圈套。特朗普现在是不讲什么外交礼仪、外交策略了，如果打个比喻，就像满地打滚，甚至拿砖头往自己头上砸。那对待这样的人，咱们最好的做法是离他远一点儿，因为你越理他，他越来劲，咱们要避免上他的圈套。'
result=ner(s)
print(result)

output：

1
2
2025-05-21 09:20:21,551 - modelscope - WARNING - Using branch: master as version is unstable, use with caution
[(14, 16, 'PERSON', '特朗普'), (36, 38, 'PERSON', '特朗普'), (58, 60, 'PERSON', '特朗普')]

另：

torchcrf包安装问题

解决方法：

1
pip install pytorch-crf

总结：

其实本质上NER任务，还是在做分类问题，只是是相对较为复杂的分类问题。