【“Transformers快速入门”学习笔记2】加载与保存模型和分词器

模型

学习笔记1中已经介绍了checkpoint方法：

from transformers import AutoModel

checkpoint = '     '
model = AutoModel.from_pretrained(checkpoint)

加载模型

所有存储在 HuggingFace Model Hub 上的模型都可以通过 Model.from_pretrained() 来加载权重，参数可以像上面一样是 checkpoint 的名称，也可以是本地路径（预先下载的模型目录），例如：

from transformers import BertModel

model = BertModel.from_pretrained("./models/bert/")

Model.from_pretrained() 会自动缓存下载的模型权重，默认保存到 ~/.cache/huggingface/transformers，我们也可以通过 HF_HOME 环境变量自定义缓存目录。

保存模型

使用 .save_pretrained 将指定模型保存到指定位置，比如：

from transformers import AutoModel

model = AutoModel.from_pretrained("bert-base-cased")
model.save_pretrained("./models/bert-base-cased/")

这会在保存路径下创建两个文件：

config.json：模型配置文件，存储模型结构参数，例如 Transformer 层数、特征空间维度等。记录模型的结构
pytorch_model.bin：又称为 state dictionary，存储模型的权重。记录模型的参数

分词器

由于神经网络模型不能直接处理文本，因此我们需要先将文本转换为数字，这个过程被称为编码 (Encoding)，其包含两个步骤：

分词，使用分词器 (tokenizer) 将文本按词、子词、字符切分为 tokens；
映射，将所有的 token 映射到对应的 token ID。

加载保存分词器

from transformers import AutoTokenizer

checkpoint = '     '
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
tokenizer.save_pretrained("./models/bert-base-cased/")

调用 Tokenizer.save_pretrained() 函数会在保存路径下创建三个文件：

special_tokens_map.json：映射文件，里面包含 unknown token 等特殊字符的映射关系；
tokenizer_config.json：分词器配置文件，存储构建分词器需要的参数；
vocab.txt：词表，一行一个 token，行号就是对应的 token ID（从 0 开始）。

编码

如上文所述，编码（Encoding）包括分词和映射两个步骤

分词，并且将tokens转换为对应的token IDs

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")

sequence = "Every tear drop is a waterfall"
tokens = tokenizer.tokenize(sequence)
ids = tokenizer.convert_tokens_to_ids(tokens)

print(tokens)
print(ids)

输出

['Every', 'tear', 'drop', 'is', 'a', 'waterfall']
[4081, 7591, 3968, 1110, 170, 16962]

以上步骤可以使用encode()合并，并且 encode() 会自动添加模型需要的特殊 token，例如 BERT 分词器会分别在序列的首尾添加 [CLS] and [SEP]

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("D:/Desktop/trans/models/bert-base-cased")

sequence = "Every tear drop is a waterfall"
sequence_ids = tokenizer.encode(sequence)

print(sequence_ids)

输出：

[101, 4081, 7591, 3968, 1110, 170, 16962, 102]

更常用的是直接使用分词器处理，可以返回更多的输出：

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("D:/Desktop/trans/models/bert-base-cased")

sequence = "Every tear drop is a waterfall"
sequence_ids = tokenizer(sequence)

print(sequence_ids)

输出：

{'input_ids': [101, 4081, 7591, 3968, 1110, 170, 16962, 102], 
'token_type_ids': [0, 0, 0, 0, 0, 0, 0, 0], 
'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1]}

解码

解码（Decoding）将token ids转换回文字。解码过程不是简单的映射回tokens，还需要合并被分为多个token的单词。

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("D:/Desktop/trans/models/bert-base-cased")

token_ids = ([101, 4081, 7591, 3968, 1110, 170, 16962, 102])
decode_string = tokenizer.decode(token_ids)

print(decode_string)

输出：

[CLS] Every tear drop is a waterfall [SEP]