【NLP 47、实践 ⑫ 通过mask和loss计算实现SFT】

目录

SFT训练

attention Mask：

代码实现

数据文件

1.模型定义

2.前向传播，计算loss

代码运行流程

3.加载语料

4.构造掩码

5.填充或截断

代码运行流程

6. SFT的数据构造

代码运行流程

7.建立模型

8.采样策略 

9.模型效果评估

10.模型训练

代码运行流程

11.完整代码 

---

别着急，怎么样都很好

                        —— 25.3.26

SFT训练

        与传统的预训练模型训练过程类似，不过要对可见的query进行Mask，也就是输入问答QA对，传统的预训练模型是全部可见，而SFT任务则需要对QA对进行Mask处理，使得答案中的每个字只能看到之前生成的答案以及全部的问题，而问题只能计算自身的注意力分数，而不能计算与答案的注意力分数

SFT过程中，我们需要一个另外形状的Mask，使其掩盖住输入中问题所对应的答案

---

attention Mask：

假定x1、x2、x3对应的是问题中的每个字，y1、y2对应的是这个问题应该给出的回答

【cls】：输入文本的句首符        【sep】：两个句子的分隔符        【sep】：输入文本的句尾符

        Transfomer架构中，会有【文本长度 × 文本长度】的矩阵代表注意力权重，在训练过程中，问题部分应该全部可见，而问题应该看不到答案；答案部分可以完全看到问题，答案中的每个字可以看到之前答案的字和所有的问题，看不到之后答案的字

        最后输出的答案部分只需计算答案部分，即输出答案y1、y2的概率分布

---

代码实现

数据文件

通过网盘分享的文件：文本生成
链接: https://pan.baidu.com/s/1Az9WLH1LfEyk_5ih8db7jw?pwd=6uv6 提取码: 6uv6 
--来自百度网盘超级会员v3的分享

---

1.模型定义

hidden_size：表示BERT模型的隐藏层维度，即每个词经过模型处理后输出的向量维度。该参数直接影响模型的表征能力

vocab_size：词汇表的大小，即模型需要预测的所有可能token的数量。

pretrain_model_path：预训练BERT模型的路径），用于加载预训练权重。预训练模型已通过大规模语料学习语言表示，微调时能显著提升下游任务性能

BertModel.from_pretrained()：加载预训练的BERT模型权重，用于生成上下文相关的词向量表示。支持从Hugging Face模型库或本地路径加载模型，适用于各类NLP任务（如文本分类、问答等）

参数	​类型	​描述
pretrained_model_name_or_path	str	预训练模型的名称（如bert-base-chinese）或本地路径。Hugging Face模型库自动下载，本地路径需包含配置文件与权重。
config	BertConfig	可选参数，自定义模型配置。若未指定，则使用默认配置。
output_attentions	bool	是否输出每层的注意力权重矩阵（形状为[batch_size, num_heads, seq_len, seq_len]）。默认False。
output_hidden_states	bool	是否输出所有隐藏层的状态（包括词嵌入层和12层Transformer）。默认False。
return_dict	bool	是否以字典形式返回结果（兼容旧版代码时设为False）。默认
attn_implementation	str	指定注意力机制的计算实现方式： - "eager"：标准PyTorch实现（兼容性好，但未优化） - "flash_attention_2"：使用Flash Attention加速计算（需硬件支持）。

nn.Linear()：定义全连接层，执行线性变换 y=xA^T+b。输入与输出均为二维张量，常用于分类器或特征转换层。

​参数	​类型	​描述
in_features	int	输入张量的特征维度（即最后一维大小）。例如，输入形状为[batch_size, 768]时，in_features=768。
out_features	int	输出张量的特征维度。例如，若需将BERT输出映射到10分类任务，则out_features=10。
bias	bool	是否启用偏置项 b。默认True。禁用时可减少参数量（适用于某些压缩场景）。
device	torch.device	指定计算设备（如cuda:0或cpu）。默认跟随模型参数。
dtype	torch.dtype	指定权重和偏置的数据类型（如torch.float16）。默认与输入张量一致。

nn.CrossEntropyLoss()：计算多分类任务的交叉熵损失，结合LogSoftmax与NLLLoss。输入为未归一化的logits，输出为标量损失值。

​参数	​类型	​描述
weight	Tensor	类别权重张量，用于处理类别不平衡（如weight=torch.tensor([0.5, 2.0])）。默认None。
ignore_index	int	指定忽略的标签索引（如填充符-1）。默认-100。
reduction	str	损失汇总方式：'none'（不汇总）、'mean'（平均）、'sum'（求和）。默认'mean'。

class LanguageModel(nn.Module):
    def __init__(self, hidden_size, vocab_size, pretrain_model_path):
        super(LanguageModel, self).__init__()
        # self.embedding = nn.Embedding(len(vocab), input_dim)
        # self.layer = nn.LSTM(input_dim, input_dim, num_layers=1, batch_first=True)

        self.bert = BertModel.from_pretrained(pretrain_model_path, return_dict=False, attn_implementation='eager')

        self.classify = nn.Linear(hidden_size, vocab_size)
        self.loss = nn.CrossEntropyLoss(ignore_index=-1)

---

2.前向传播，计算loss

代码运行流程

forward 方法执行流程
├── 输入参数
│   ├── x: 输入数据（如文本序列的token_id）
│   ├── mask: 注意力掩码矩阵（训练时控制交互范围）
│   └── y: 真实标签（存在时为训练模式，否则为预测模式）
│
├── 模式分支判断
│   └── if y is not None → 训练模式
│       └── else → 预测模式
│
├── 训练模式（计算Loss）
│   ├── 1. 注意力掩码处理
│   │   └── print(mask.shape) → 验证掩码形状（调试用）
│   ├── 2. BERT模型前向传播
│   │   ├── 调用 self.bert(x, attention_mask=mask)
│   │   │   └── 递归调用 BERT 的 __call__ 方法 → 触发其 forward
│   │   └── 返回特征向量 x 和池化输出（_ 被忽略）
│   ├── 3. 分类层预测
│   │   └── y_pred = self.classify(x) → 形状 (batch_size, vocab_size)
│   └── 4. Loss计算
│       ├── 维度展平：y_pred.view(-1, vocab_size) → 适配交叉熵输入
│       └── self.loss(y_pred, y.view(-1)) → 计算预测与标签的损失值
│
├── 预测模式（生成概率）
│   ├── 1. BERT模型前向传播
│   │   ├── 调用 self.bert(x) → 无掩码控制
│   │   │   └── 默认使用全连接注意力（允许所有位置交互）
│   │   └── 返回特征向量 x 和池化输出
│   ├── 2. 分类层预测
│   │   └── y_pred = self.classify(x) → 形状 (batch_size, vocab_size)
│   └── 3. 概率归一化
│       └── torch.softmax(y_pred, dim=-1) → 输出各类别概率分布
│
└── PyTorch底层机制（基于搜索结果
    ├── 隐式调用逻辑
    │   └── 外部调用 model(x, mask, y) → 触发 __call__ 方法 → 执行 forward
    └── 子模块递归处理
        └── self.bert 和 self.classify 均为 nn.Module 子类，其 forward 被递归调用

x：输入序列的Token ID矩阵（张量），形状为 (batch_size, sequence_length)

mask：注意力掩码矩阵，形状为 (batch_size, sequence_length)

过self.bert(x)后的x： BERT 最后一层的隐藏状态 last_hidden_state，形状为 (batch_size, seq_length, hidden_size)

过self.bert(x)后的第二个输出_： 忽略 pooler_output（用于分类任务的 [CLS] 向量）

y：真实标签张量，形状与任务相关（如分类任务为 (batch_size,)，语言模型任务为 (batch_size, sequence_length)

y_pred：计算出的预测标签张量，形状为 (batch_size, seq_length, vocab_size)（通过全连接层映射到词表大小）

view()：调整张量形状，不改变数据存储顺序，要求张量内存连续（类似 reshape 但更严格）

​参数	​类型	​描述
*shape	int 或 tuple	目标形状，支持动态推断（如 view(-1, 4) 自动计算行数）
​返回值	torch.Tensor	共享数据存储的新视图张量，与原张量数据同步修改

shape()：返回张量/数组的维度元组，用于获取多维数据的结构信息

self.bert()：加载预训练的BERT模型，执行文本编码并输出隐藏状态

​参数	​类型	​描述
input_ids	torch.Tensor	​必需，输入Token ID矩阵，形状 (batch_size, seq_length)
attention_mask	torch.Tensor	​可选，掩码矩阵（1有效/0无效），用于忽略填充符
output_attentions	bool	​可选，是否返回所有注意力权重矩阵（默认 False）
output_hidden_states	bool	​可选，是否返回所有隐藏层输出（默认 False）
​返回值	tuple 或 dict	包含 last_hidden_state, pooler_output 等

torch.softmax()：将Logits转换为概率分布，确保各维度概率和为1，用于多分类任务

参数	​类型	​描述
input	torch.Tensor	​必需，输入Logits张量（如分类层输出）
dim	int	​必需，指定计算维度（如 dim=-1 对最后一个维度归一化）
​返回值	torch.Tensor	概率分布张量，各指定维度元素和为1

    # 当输入真实标签，返回loss值；无真实标签，返回预测值
    def forward(self, x, mask=None, y=None):
        if y is not None:
            # 训练时，构建一个下三角的mask矩阵，让上下文之间没有交互
            print(mask.shape)
            x, _ = self.bert(x, attention_mask=mask)
            y_pred = self.classify(x)  # output shape:(batch_size, vocab_size)
            return self.loss(y_pred.view(-1, y_pred.shape[-1]), y.view(-1))
        else:
            # 预测时，可以不使用mask
            x, _ = self.bert(x)
            y_pred = self.classify(x)  # output shape:(batch_size, vocab_size)
            return torch.softmax(y_pred, dim=-1)

---

3.加载语料

        加载语料，因为文本数据是title和content的标签对，所以将title部分当成假想提示词prompt，将content部分当成假想的答案answer

path：​输入文件路径，指向存储语料数据的JSON文件（每行一个JSON对象）

corpus：输出列表，存储处理后的语料数据，每个元素为 [title, content] 的二元组

f：​文件对象，以UTF-8编码打开文件，用于逐行读取内容

line：​单行文本，对应文件中的一行原始JSON字符串（需解析为字典）

open()：打开文件并返回文件对象，用于文件的读取、写入等操作

​参数	​类型	​描述	​示例
file	str	​必需，文件路径（绝对或相对路径）	open("data.txt", "r")
mode	str	​可选，文件打开模式（默认'r'，文本只读）	'rb'（二进制只读）
buffering	int	​可选，缓冲区大小（默认-1，系统自动优化）	buffering=0（无缓冲）
encoding	str	​可选，文件编码（默认None，系统编码）	encoding='utf-8'
errors	str	​可选，编码错误处理方式（如'ignore'忽略错误）	errors='replace'
newline	str	​可选，换行符控制（默认None，自动识别）	newline='\n'
​常见模式	-	'r'（只读）、'w'（覆盖写入）、'a'（追加）、'b'（二进制模式）	'a+'（追加读写）

json.loads()：将JSON格式的字符串解析为Python对象（如字典、列表等）

​参数	​类型	​描述	​示例
json_str	str	​必需，需要解析的JSON字符串	json.loads('{"name": "Alice"}')
object_hook	Callable	​可选，自定义解析函数（将JSON对象转换为特定Python对象）	object_hook=lambda d: CustomClass(**d)
parse_float	Callable	​可选，自定义浮点数解析方式（如转为Decimal）	parse_float=decimal.Decimal
parse_int	Callable	​可选，自定义整数解析方式（如转为十六进制）	parse_int=lambda x: int(x, 16)
parse_constant	Callable	​可选，处理特殊常量（如NaN、Infinity）	-
object_pairs_hook	Callable	​可选，自定义键值对解析函数（替代默认字典）

列表.append()：在列表末尾添加单个元素，直接修改原列表

​参数	​类型	​描述	​示例
element	Any	​必需，要添加到列表末尾的元素（可以是任意类型）	list1.append(42)
​返回值	-	无返回值，直接修改原列表	list1 = [1,2]; list1.append(3) → [1,2,3]

# 加载语料, 用title当成假想的prompt，content当成假想的answer
def load_corpus(path):
    corpus = []
    with open(path, encoding="utf8") as f:
        for line in f:
            line = json.loads(line)
            corpus.append([line["title"], line["content"]])
    return corpus

---

4.构造掩码

s1：第一个字符串的原始长度​（不含特殊标记），如输入句子1的Token数

s2：第二个字符串的原始长度​（不含特殊标记），如输入句子2或回答的Token数

len_s1：​扩展后的第一个序列长度，包含 [CLS] 和 [SEP] 标记

len_s2：扩展后的第二个序列长度，包含 [SEP] 标记

mask：注意力掩码矩阵，控制Token之间的可见性

i：​循环索引，用于遍历Token位置

torch.ones()：创建指定形状的全1张量，支持自定义数据类型、设备等属性

参数	​类型	​描述	​示例
*size	int 或 tuple	​必需，张量的形状（如 (2,3) 或 2,3）	torch.ones(2, 3)
dtype	torch.dtype	​可选，张量数据类型（默认与全局默认类型一致）	dtype=torch.float32
layout	torch.layout	​可选，内存布局（默认torch.strided）	layout=torch.sparse_coo
device	torch.device	​可选，张量存储设备（默认当前设备）	device='cuda'
requires_grad	bool	​可选，是否启用自动微分（默认False）	requires_grad=True

range()：生成一个不可变的整数序列，常用于循环控制或生成索引序列

参数	​类型	​描述	​示例
start	int	​可选，序列起始值（默认0）	range(2, 5) → 2,3,4
stop	int	​必需，序列结束值（不包含该值本身）	range(3) → 0,1,2
step	int	​可选，步长（默认1，可为负数）	range(0, 10, 2) → 0,2,4,6,8

# 构造掩码，输入两个字符串的长度
def create_mask(s1, s2):
    len_s1 = s1 + 2  # cls + sep
    len_s2 = s2 + 1  # sep
    # 创建掩码张量
    mask = torch.ones(len_s1 + len_s2, len_s1 + len_s2)
    # 遍历s1的每个token
    for i in range(len_s1):
        # s1的当前token不能看到s2的任何token
        # ​行索引：i（取值范围：0 ≤ i < len_s1）
        # 对应第一个序列（s1）的每个Token位置（包含[CLS]和[SEP]标记）。
        # ​列索引：len_s1:（从len_s1到末尾）
        # 对应第二个序列（s2）的所有Token位置（包含s2的[SEP]）。
        mask[i, len_s1:] = 0
        
    for i in range(len_s2):
        # 遍历s2的每个token
        # ​行索引：len_s1 + i（取值范围：len_s1 ≤ lens1 + i < len_s1 + len_s2）
        # 对应第二个序列（s2）的每个Token位置（包含[SEP]标记）。
        # ​列索引：len_s1 + i + 1:（从len_s1 + i + 1到末尾）
        # 对应当前Token位置之后的所有位置（包括后续的s2的Token和s1的Token)
        # s2的当前token不能看到后面的s2 token
        mask[len_s1 + i, len_s1 + i + 1:] = 0
    return mask

---

5.填充或截断

代码运行流程

pad_mask 函数运行流程
├── 输入参数
│   ├── tensor: 输入二维张量（如序列矩阵） → 形状 (height, width)
│   └── target_shape: 目标尺寸 → (target_height, target_width)
│
├── 初始化阶段
│   ├── 获取原始尺寸: height, width = tensor.shape
│   └── 创建全零结果张量:
│       └── result = torch.zeros(target_shape, dtype, device) → 兼容原设备与类型
│
├── 填充/截断逻辑
│   ├── 左上角对齐策略:
│       ├── h_start = w_start = 0 → 原始数据始终位于结果左上角
│       ├── h_end = min(height, target_height) → 动态适配高度
│       └── w_end = min(width, target_width) → 动态适配宽度
│   └── 数据复制:
│       └── result[0:h_end, 0:w_end] = tensor[:h_end, :w_end] → 截断或保留原数据
│
├── 输出结果
│   └── return result → 形状为 target_shape 的数值填充张量
│
└── 下游应用（与掩码生成关联)
    ├── 布尔掩码生成:
        └── mask = (result != 0) → 标记有效数据位置（True=有效，False=填充）
    └── 注意力机制适配:
        ├── 维度扩展 → 适配多头注意力头数（如 [batch_size, num_heads, seq_len, seq_len]）
        └── 掩码作用 → 遮蔽填充符的注意力权重（Softmax 前替换为 -inf）

tensor：​必需参数，二维输入张量，通常是需调整形状的原始数据（如注意力分数矩阵或特征图）

target_shape：​必需参数，目标形状 (target_height, target_width)，表示输出张量的尺寸。

height，weight：原始输入张量的原始高度和宽度，通过 tensor.shape 获取。

target_height，target_weight：目标张量的高度和宽度，来自 target_shape

result：初始化全零张量，形状为 target_shape，用于存储填充或截断后的数据。

h_start，w_start：填充起始位置（始终为0，表示从左上角开始填充原始数据）。

h_end，w_end：填充或截断的终止位置，取原始尺寸与目标尺寸的最小值，防止越界。

torch.zeros()：创建指定形状的全零张量，用于初始化权重、占位符或存储结构化数据

参数	​类型	​描述	​示例
*size	int 或 tuple	​必需，定义张量的形状（如 3 或 (2,3)）	torch.zeros(2,3) → 2x3全零矩阵
dtype	torch.dtype	​可选，指定张量数据类型（默认 torch.float32）	dtype=torch.int64
device	torch.device	​可选，指定存储设备（如 'cpu' 或 'cuda'，默认跟随全局设置）	device='cuda'
requires_grad	bool	​可选，是否启用梯度计算（默认 False）

min()：返回可迭代对象或多个参数中的最小值，支持自定义比较逻辑（如通过 key 参数）和处理空输入（如 default 参数）

​参数	​类型	​描述	​示例
iterable	可迭代对象	​必需​（单参数形式），需比较的列表、元组等	min([3,1,4]) → 1
arg1, arg2, ...	任意类型	​必需​（多参数形式），直接比较多个值	min(3, 1, 4) → 1
key	callable	​可选，自定义比较函数（如 key=lambda x: len(x) 找最短字符串）	key=abs（按绝对值比较）
default	任意类型	​可选，当 iterable 为空时的返回值（不指定则抛 ValueError）	default=0（空列表返回0）

张量.shape：返回张量的维度信息（形状），以元组形式表示各维度长度，用于维度校验、重塑操作或动态调整计算流程

属性/方法	​类型	​描述	​示例
​返回值	torch.Size	张量形状的元组（如 (2,3) 表示2行3列）	x.shape → torch.Size([2,3])
​关联操作	-	常用配合方法： - .ndim：维度数（如2D张量的 ndim 为2） - .numel()：元素总数	x.ndim → 2 x.numel() → 6

def pad_mask(tensor, target_shape):
    # 获取输入张量和目标形状的长宽
    height, width = tensor.shape
    target_height, target_width = target_shape
    # 创建一个全零张量,形状为目标形状
    result = torch.zeros(target_shape, dtype=tensor.dtype, device=tensor.device)
    # 计算需要填充或截断的区域
    h_start = 0
    w_start = 0
    h_end = min(height, target_height)
    w_end = min(width, target_width)
    # 将原始张量对应的部分填充到全零张量中
    result[h_start:h_end, w_start:w_end] = tensor[:h_end - h_start, :w_end - w_start]
    return result

---

6. SFT的数据构造

代码运行流程

build_dataset 函数运行流程
├── 1. 输入参数
│   ├── tokenizer: 分词器，用于文本编码
│   ├── corpus: 原始语料库，格式为 (prompt, answer) 列表
│   ├── max_length: 最大序列长度（控制填充/截断）
│   └── batch_size: 数据加载的批次大小
│
├── 2. 遍历语料库
│   └── 对每个 (prompt, answer) 对执行以下操作：
│       ├── 2.1 文本编码
│       │   ├── prompt_encode: 分词器编码prompt（无特殊标记）
│       │   └── answer_encode: 分词器编码answer（无特殊标记）
│       │
│       ├── 2.2 序列构造
│       │   ├── x = [CLS] + prompt + [SEP] + answer + [SEP]
│       │   └── y = (prompt部分全为-1) + answer + [SEP] + 末尾填充-1
│       │       │   - prompt部分标记为-1：不参与loss计算
│       │       │   - answer部分保留原始token_id：参与loss计算
│       │
│       ├── 2.3 掩码矩阵生成
│       │   └── mask = create_mask(prompt_len, answer_len)
│       │       │   - prompt内部允许双向注意力（可交互）
│       │       │   - answer部分仅允许自回归注意力（单向交互）
│       │
│       ├── 2.4 填充处理
│       │   ├── x: 截断至max_length，右侧补0
│       │   ├── y: 截断至max_length，右侧补0
│       │   └── mask: 调用pad_mask扩展至(max_length, max_length)
│       │
│       └── 2.5 张量转换
│           ├── x: 转为LongTensor
│           ├── y: 转为LongTensor
│           └── mask: 保持与输入一致的dtype
│
├── 3. 数据封装
│   └── 返回DataLoader
│       ├── 数据集: 包含[x, mask, y]的列表
│       ├── batch_size: 控制批次大小
│       └── shuffle=True: 打乱数据顺序
│
└── 4. 下游训练关联
    ├── 4.1 注意力控制
    │   └── mask确保：
    │       ├── prompt内部全连接（允许自由交互）
    │       └── answer部分仅允许自回归注意力（防止信息泄露）
    │
    └── 4.2 Loss计算
        ├── 仅y中非-1位置参与计算
        └── 通过mask限制注意力范围，实现自回归生成

tokenizer：分词器，用于将文本转换为Token ID序列

corpus：训练语料，包含(prompt, answer)对

max_length：​最大序列长度，控制输入截断和填充

batch_size：​批处理大小，决定每次训练迭代的样本数

dataset：存储处理后的训练样本集合，包含输入序列、注意力掩码和标签，用于后续的批处理训练

• 每个样本由三部分构成：
  • ​x：输入序列（含特殊标记），例如 [CLS] + prompt + [SEP] + answer + [SEP]
  • ​mask：注意力掩码矩阵，控制序列内Token的可见性（如prompt内部双向可见，answer部分因果掩码）
  • ​y：标签序列，仅对answer部分计算损失（通过-1屏蔽其他区域）
• 最终通过 DataLoader 封装为可迭代的批处理数据，支持并行加载与训练

prompt：输入的问题或指令部分，用于引导模型生成特定类型的回答。

answer：模型需要生成的回答部分，标签中仅此部分参与损失计算。

prompt_encode：​编码后的prompt，不含特殊标记

answer_encode：​编码后的answer，不含特殊标记

tokenizer.cls_token_id：表示分类标记（[CLS]）的 Token ID，用于标识序列的起始位置，并为模型提供全局语义聚合的锚点

tokenizer.sep_token_id：表示分隔符（[SEP]）的 Token ID，用于划分输入中的不同片段（如 prompt 与 answer），或在序列结尾标识终止

x：完整输入序列，包含特殊标记

y：​标签序列，仅对answer部分计算损失

mask：注意力掩码矩阵，控制Token可见性

enumerate()：将可迭代对象（列表、元组、字符串等）转换为索引序列，返回由 (索引, 元素) 组成的元组迭代器

​参数	​类型	​描述	​示例
iterable	可迭代对象	必需参数，需遍历的对象（如列表、字符串）	enumerate(['a', 'b'])
start	int	可选参数，索引起始值（默认 0）	enumerate(['a', 'b'], start=1)

tokenizer.encode()：将可迭代对象（列表、元组、字符串等）转换为索引序列，返回由 (索引, 元素) 组成的元组迭代器

参数	​类型	​描述	​示例
text	str 或 List	必需参数，需编码的文本或分词后的列表	"Hello, world!"
add_special_tokens	bool	可选参数，是否添加特殊标记（如 [CLS]、[SEP]），默认 True	add_special_tokens=False
max_length	int	可选参数，最大序列长度，超出部分截断	max_length=512
padding	bool 或 str	可选参数，填充策略（如 True 填充至最长序列，'max_length' 填充至指定长度）	padding='max_length'
truncation	bool 或 str	可选参数，截断策略（如 True 自动截断，'only_first' 仅截断首句）	truncation=True

torch.LongTensor()：创建 64 位整数类型的张量，用于存储整型数据（如索引、标签）

​参数	​类型	​描述	​示例
data	列表/数组/张量	必需参数，整数型数据（如列表、NumPy 数组）	torch.LongTensor([1, 2, 3])

列表.append()：在列表末尾添加单个元素，直接修改原列表，无返回值

​参数	​类型	​描述	​示例
obj	任意类型	必需参数，需添加到列表末尾的元素（支持整数、字符串、列表等）	list1.append(5)

DataLoader()：将数据集封装为可迭代的批处理对象，支持自动分批、打乱顺序和多线程加载

​参数	​类型	​描述	​示例
dataset	Dataset 实例	必需参数，自定义数据集对象（需实现 __len__ 和 __getitem__ 方法）	DataLoader(mnist_dataset)
batch_size	int	可选参数，每批样本数（默认 1）	batch_size=32
shuffle	bool	可选参数，是否打乱数据顺序（默认 False，验证集通常关闭）	shuffle=True
num_workers	int	可选参数，数据加载的并行进程数（0 表示主进程加载）	num_workers=4
drop_last	bool	可选参数，是否丢弃最后不完整的批次（默认 False）	drop_last=True
collate_fn	callable	可选参数，自定义批次合并逻辑（如动态填充）

# sft的数据构造
# loss只计算答案部分，通过mask矩阵，让上下文之间没有交互
# label中使用-1，表示不参与训练
def build_dataset(tokenizer, corpus, max_length, batch_size):
    dataset = []
    for i, (prompt, answer) in enumerate(corpus):
        prompt_encode = tokenizer.encode(prompt, add_special_tokens=False)
        answer_encode = tokenizer.encode(answer, add_special_tokens=False)
        x = [tokenizer.cls_token_id] + prompt_encode + [tokenizer.sep_token_id] + answer_encode + [
            tokenizer.sep_token_id]
        y = len(prompt_encode) * [-1] + [-1] + answer_encode + [tokenizer.sep_token_id] + [-1]
        # 构建一个的mask矩阵，让prompt内可以交互，answer中上下文之间没有交互
        mask = create_mask(len(prompt_encode), len(answer_encode))
        # padding
        x = x[:max_length] + [0] * (max_length - len(x))
        y = y[:max_length] + [0] * (max_length - len(y))
        x = torch.LongTensor(x)
        y = torch.LongTensor(y)
        mask = pad_mask(mask, (max_length, max_length))
        dataset.append([x, mask, y])

    return DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=0)

---

7.建立模型

vocab：定义模型的词汇表大小或具体词汇列表。

char_dim：指定字符嵌入向量的维度，即每个词元在模型中的向量表示长度。

pretrain_model_path：预训练模型权重的加载路径，用于迁移学习或微调任务。

768：对应 char_dim，定义词向量维度

21128：对应 vocab 的词汇表大小

# 建立模型
def build_model(vocab, char_dim, pretrain_model_path):
    model = LanguageModel(768, 21128, pretrain_model_path)
    return model

---

8.采样策略 

prob_distribution：模型输出的概率分布（需归一化），表示每个词被选中的概率。维度为 [vocab_size]。

strategy：根据随机数选择策略："greedy" 或 "sampling"。

random.random()：生成 [0.0, 1.0) 之间的随机浮点数，符合均匀分布。常用于概率判断或生成随机测试数据

torch.argmax()：返回张量中最大值所在的索引，支持按维度计算或全局索引

参数	​类型	​描述	​示例
input	torch.Tensor	必需参数，输入张量	torch.tensor([1, 3, 2])
dim	int 或 None	可选参数，指定计算维度（默认 None 表示全局索引）	dim=1（按行计算）
keepdim	bool	可选参数，是否保持原维度（默认 False）	keepdim=True（输出维度不变）

tensor.cpu()：将张量从 GPU 转移到 CPU 内存，便于与 NumPy 等非 GPU 库交互

tensor.numpy()：将 PyTorch 张量转换为 NumPy 数组，实现与 Python 生态的无缝交互

np.random.choice()：从数组或整数范围中随机抽取元素，支持权重和重复抽样

​参数	​类型	​描述	​示例
a	一维数组或整数	必需参数，输入数据源（整数时等价于 np.arange(a)）	a=[1, 2, 3] 或 a=5
size	int 或 tuple	可选参数，输出形状（默认 None 返回单个值）	size=3（抽取3个元素）
replace	bool	可选参数，是否允许重复抽样（默认 True）	replace=False（无放回抽样）
p	一维数组	可选参数，每个元素的抽样概率（默认均匀分布）	p=[0.1, 0.3, 0.6]

list()：将可迭代对象（如元组、字符串、集合）转换为列表

​参数	​类型	​描述	​示例
iterable	可迭代对象	可选参数，输入数据（默认生成空列表）	list((1, 2, 3)) → [1, 2, 3]

range()：生成不可变的整数序列，常用于循环控制或列表推导式

​参数	​类型	​描述	​示例
start	int	可选参数，起始值（默认 0）	range(5) → 0,1,2,3,4
stop	int	必需参数，终止值（不包含该值）	range(2, 5) → 2,3,4
step	int	可选参数，步长（默认 1，支持负数逆向生成）	range(0, 10, 2) → 0,2,4,6,8

# 采样策略选择
def sampling_strategy(prob_distribution):
    if random.random() > 0.1:
        strategy = "greedy"
    else:
        strategy = "sampling"
    if strategy == "greedy":
        return int(torch.argmax(prob_distribution))
    elif strategy == "sampling":
        prob_distribution = prob_distribution.cpu().numpy()
        return np.random.choice(list(range(len(prob_distribution))), p=prob_distribution)

---

9.模型效果评估

openings：​必需参数，初始输入文本，用于引导生成方向（如开头句）

model：​必需参数，预训练的语言模型，用于预测下一个词的概率分布

tokenizer：必需参数，将文本与词ID相互转换的工具，支持编码（encode）和解码（decode）。

x：将词ID列表转换为张量输入模型，支持GPU加速（.cuda()）。

y：模型输出的最后一个词的概率分布，维度为词表大小。

index：通过采样策略选出的下一个词ID，添加到生成序列中。

model.eval()：将模型切换为评估模式，关闭训练相关层（如 Dropout 和 Batch Normalization 的训练行为），确保推理时参数稳定

tokenizer.encode()：将文本编码为词 ID 序列，支持添加特殊标记（如 [CLS]、[SEP]）和填充对齐

​参数	​类型	​描述	​示例
text	str 或 List[str]	​必需，输入文本（单句或句对）	"自然语言处理"
add_special_tokens	bool	​可选，是否添加特殊标记（默认 True）	add_special_tokens=False
max_length	int	​可选，最大序列长度（超过则截断）	max_length=128
padding	str	​可选，填充策略（"max_length"、"longest" 或 False，默认 False）	padding="max_length"
truncation	bool 或 str	​可选，是否截断超长文本（默认 False）	truncation=True
return_tensors	str	​可选，返回张量类型（如 "pt" 返回 PyTorch 张量，默认不返回）	return_tensors="pt"

torch.no_grad()：禁用梯度计算，减少内存消耗并加速推理过程

torch.LongTensor()：创建 64 位整数类型的张量，常用于存储词 ID 或索引

​参数	​类型	​描述	​示例
data	list 或 int	​必需，输入数据（列表或标量）	torch.LongTensor([101, 2769])

torch.cuda.is_available()：检查当前环境是否支持 CUDA（GPU 加速）

cuda()：PyTorch 中用于将张量从 CPU 迁移到 GPU 的函数，以利用 GPU 并行计算加速张量运算。该函数会返回一个新的 GPU 张量副本，原张量仍保留在 CPU 上（非原地操作）。通过显存管理优化，可显著提升深度学习模型的训练和推理效率

​参数名	​类型	​描述	​示例
device	int 或 str	​可选参数，指定目标 GPU 设备（默认 cuda:0）	0 或 "cuda:0"
non_blocking	bool	​可选参数，是否启用异步传输（默认 False，需与 pin_memory=True 配合）

列表.append()：向列表末尾添加元素，常用于动态扩展序列

参数	​类型	​描述	​示例
element	任意类型	​必需，要添加的元素	tokens.append(102)

tokenizer.decode()：将词 ID 序列解码为文本，支持跳过特殊标记

​参数	​类型	​描述	​示例
token_ids	List[int]	​必需，输入词 ID 序列	[101, 2769, 102]
skip_special_tokens	bool	​可选，是否跳过特殊标记（如 [CLS]、[SEP]，默认 False）	skip_special_tokens=True
clean_up_tokenization_spaces	bool	​可选，是否清理多余空格（默认 True）

# 文本生成测试代码
def generate_sentence(openings, model, tokenizer):
    model.eval()
    openings = tokenizer.encode(openings)
    with torch.no_grad():
        # 生成文本超过50字则终止迭代
        while len(openings) <= 50:
            x = torch.LongTensor([openings])
            if torch.cuda.is_available():
                x = x.cuda()
            y = model(x)[0][-1]
            index = sampling_strategy(y)
            openings.append(index)
    return tokenizer.decode(openings)

---

10.模型训练

代码运行流程

main 函数执行流程
├── 1. 参数初始化阶段
│   ├── 超参数设定
│   │   ├── epoch_num=20（训练轮次）
│   │   ├── batch_size=32（批处理量）
│   │   ├── char_dim=768（BERT 特征维度）
│   │   └── learning_rate=0.001（学习率）
│   │
│   └── 预训练模型加载
│       └── tokenizer = BertTokenizer.from_pretrained() → 加载BERT分词器
│
├── 2. 数据准备阶段
│   ├── 语料加载
│   │   └── corpus = load_corpus(corpus_path) → 原始文本读取
│   │
│   └── 数据集构建
│       └── train_data = build_dataset() → 生成带掩码的批次数据
│           │   - 包含编码后的 x、注意力掩码 mask 和标签 y
│
├── 3. 模型构建阶段
│   ├── 模型初始化
│   │   └── model = build_model(vocab_size, char_dim) → 包含BERT和分类层
│   │       │   - 继承 nn.Module 的 forward 方法（如网页1中前向计算定义）
│   │
│   └── GPU加速配置
│       └── model.cuda() → 启用CUDA计算（若可用）
│
├── 4. 优化器配置
│   └── optim = Adam(model.parameters(), lr=0.001) → 参数更新算法
│
├── 5. 训练循环（核心）
│   └── for epoch in range(20):
│       ├── 5.1 训练模式激活
│       │   └── model.train() → 开启梯度计算
│       │
│       ├── 5.2 批次数据处理
│       │   └── for x, mask, y in train_data:
│       │       ├── GPU迁移：x/mask/y.cuda() → 数据载入显存
│       │       ├── 梯度清零：optim.zero_grad() → 防止梯度累积
│       │       ├── 前向传播：loss = model(x, mask, y) → 计算交叉熵损失
│       │       ├── 反向传播：loss.backward() → 计算参数梯度
│       │       └── 参数更新：optim.step() → 更新模型权重
│       │
│       ├── 5.3 训练监控
│       │   ├── watch_loss 记录每轮平均损失
│       │   └── 打印生成样例（如"北京明年拟推..."续写结果）
│       │
│       └── 5.4 结果输出
│           └── print(f"第{epoch}轮平均loss:{np.mean(watch_loss)}")
│
└── 6. 模型保存阶段
    └── if save_weight:
        ├── 路径处理：base_name.replace("txt","pth")
        └── torch.save(model.state_dict(), model_path) → 保存权重文件

corpus_path：​必需，语料文件路径（如 data/news.txt）

save_weight：可选，是否保存训练后的模型权重（默认 True）

epoch_num：训练总轮数，控制模型学习数据集的次数

batch_size：每批训练样本数量，影响内存占用和梯度稳定性

char_dim：字符嵌入维度，决定每个字符的向量表示长度

max_length：输入序列的最大长度，超出部分截断，不足则填充

vocab_size：词汇表大小，需与预训练模型一致（BERT-base中文为21128）

learning_rate：优化器的学习率，控制参数更新步长

pretrain_model_path：预训练模型路径，用于加载分词器和模型权重

tokenizer：分词器，将文本编码为词ID序列

corpus：加载后的语料数据，包含 (prompt, answer) 对

train_data：数据集加载器，提供批量数据迭代

model：语言模型，基于预训练BERT结构构建

optim：优化器，更新模型参数以减少损失

watch_loss：记录每批次训练的损失值，用于计算轮次平均损失

loss：训练的损失值

x：输入词ID序列，形状为 (batch_size, seq_len)

y：标签序列，仅答案部分有效（非 -1），形状为 (batch_size, seq_len)

mask：注意力掩码矩阵，控制模型对无效位置的关注

base_name：用于从语料文件路径中提取基础文件名，并将其扩展名从 .txt 转换为 .pth，生成模型权重文件的名称。

model_path：​定义模型权重文件的完整保存路径，将生成的 base_name 与目标目录（如 model/）结合。

BertModel.from_pretrained()：加载预训练的BERT模型权重，用于生成上下文相关的词向量表示。支持从Hugging Face模型库或本地路径加载模型，适用于各类NLP任务（如文本分类、问答等）

参数	​类型	​描述
pretrained_model_name_or_path	str	预训练模型的名称（如bert-base-chinese）或本地路径。Hugging Face模型库自动下载，本地路径需包含配置文件与权重。
config	BertConfig	可选参数，自定义模型配置。若未指定，则使用默认配置。
output_attentions	bool	是否输出每层的注意力权重矩阵（形状为[batch_size, num_heads, seq_len, seq_len]）。默认False。
output_hidden_states	bool	是否输出所有隐藏层的状态（包括词嵌入层和12层Transformer）。默认False。
return_dict	bool	是否以字典形式返回结果（兼容旧版代码时设为False）。默认
attn_implementation	str	指定注意力机制的计算实现方式： - "eager"：标准PyTorch实现（兼容性好，但未优化） - "flash_attention_2"：使用Flash Attention加速计算（需硬件支持）。

torch.cuda.is_available()：检测当前环境是否支持 CUDA（即 GPU 是否可用）

cuda()：将张量从 CPU 迁移到 GPU，加速计算

​参数	​类型	​描述	​示例
device	int 或 str	​可选，目标 GPU 设备（如 0 或 "cuda:0"）	x = x.cuda(device=0)
non_blocking	bool	是否异步传输（需配合 pin_memory=True 使用）	non_blocking=True

torch.optim.Adam()：定义 Adam 优化器，结合动量与自适应学习率，用于更新模型参数

​参数	​类型	​描述	​默认值/示例
params	iterable	​必需，待优化的参数（如 model.parameters()）	model.parameters()
lr	float	学习率（控制参数更新步长）	0.001
betas	Tuple[float, float]	一阶和二阶动量衰减系数	(0.9, 0.999)
eps	float	数值稳定性分母修正项	1e-8
weight_decay	float	L2 正则化系数（防止过拟合）	0

model.parameters()：返回模型所有可训练参数的迭代器（用于优化器初始化）

model.train()：切换模型为训练模式，启用 Dropout 和 Batch Normalization 的训练行为

optim.zero_grad()：清空模型参数的梯度，避免梯度累积

loss.backward()：反向传播计算梯度（链式法则），梯度存储在张量的 .grad 属性中

​参数	​类型	​描述	​示例
retain_graph	bool	是否保留计算图（用于多次反向传播）	loss.backward(retain_graph=True)

optim.step()：根据梯度更新模型参数（需在 loss.backward() 后调用）

列表.append()：向列表末尾添加元素，动态扩展序列

​参数	​类型	​描述	​示例
element	任意类型	​必需，要添加的元素	tokens.append(102)

np.mean()：计算数组元素的平均值（支持多维数组和指定轴计算）

​参数	​类型	​描述	​示例
a	array_like	​必需，输入数组	np.mean([1, 2, 3])
axis	int 或 None	计算平均值的轴（None 表示全局平均）	axis=0（沿列计算）

os.path.basename()：从文件路径中提取文件名（忽略目录部分）

​参数	​类型	​描述	​示例
path	str	​必需，输入文件路径	os.path.basename("data/test.txt") → "test.txt"

os.path.join()：拼接多个路径组件，自动处理操作系统差异（如 / 与 \）

参数	​类型	​描述	​示例
*paths	str	​可变参数，路径组件列表	os.path.join("model", "weights.pth") → "model/weights.pth"

torch.save()：保存模型参数或张量到磁盘文件（支持 .pth 格式）

​参数	​类型	​描述	​示例
obj	object	​必需，要保存的对象（如 model.state_dict()）	torch.save(model.state_dict(), "model.pth")
f	str 或 文件对象	保存路径或文件对象	"model.pth"

model.state_dict()：返回模型参数的字典（键为参数名，值为 torch.Tensor），用于保存或加载模型

def main(corpus_path, save_weight=True):
    epoch_num = 20  # 训练轮数
    batch_size = 32  # 每次训练样本个数
    char_dim = 768  # 每个字的维度
    max_length = 50  # 样本文本长度
    vocab_size = 21128  # 字表大小
    learning_rate = 0.001  # 学习率

    pretrain_model_path = r"F:\人工智能NLP\NLP资料\week6 语言模型\bert-base-chinese"
    tokenizer = BertTokenizer.from_pretrained(pretrain_model_path)

    corpus = load_corpus(corpus_path)  # 加载语料
    train_data = build_dataset(tokenizer, corpus, max_length, batch_size)  # 建立数据集
    model = build_model(vocab_size, char_dim, pretrain_model_path)  # 建立模型
    if torch.cuda.is_available():
        model = model.cuda()
    optim = torch.optim.Adam(model.parameters(), lr=learning_rate)  # 建立优化器
    print("文本词表模型加载完毕，开始训练")
    for epoch in range(epoch_num):
        model.train()
        watch_loss = []
        for x, mask, y in train_data:  # 构建一组训练样本
            if torch.cuda.is_available():
                x, mask, y = x.cuda(), mask.cuda(), y.cuda()
            optim.zero_grad()  # 梯度归零
            loss = model(x, mask, y)  # 计算loss
            loss.backward()  # 计算梯度
            optim.step()  # 更新权重
            watch_loss.append(loss.item())
        print("=========\n第%d轮平均loss:%f" % (epoch + 1, np.mean(watch_loss)))
        print(generate_sentence("北京明年拟推工作日半价观看电影", model, tokenizer))
        print(generate_sentence("南京一合金厂锅炉发生爆炸", model, tokenizer))
    if not save_weight:
        return
    else:
        base_name = os.path.basename(corpus_path).replace("txt", "pth")
        model_path = os.path.join("model", base_name)
        torch.save(model.state_dict(), model_path)
        return

---

11.完整代码 

# coding:utf8

import json
import torch
import torch.nn as nn
import numpy as np
import math
import random
import os
import re
from transformers import BertTokenizer, BertModel
from torch.utils.data import Dataset, DataLoader

"""
基于Bert结构，进行sft形式的训练
"""


class LanguageModel(nn.Module):
    def __init__(self, hidden_size, vocab_size, pretrain_model_path):
        super(LanguageModel, self).__init__()
        # self.embedding = nn.Embedding(len(vocab), input_dim)
        # self.layer = nn.LSTM(input_dim, input_dim, num_layers=1, batch_first=True)

        self.bert = BertModel.from_pretrained(pretrain_model_path, return_dict=False, attn_implementation='eager')

        self.classify = nn.Linear(hidden_size, vocab_size)
        self.loss = nn.CrossEntropyLoss(ignore_index=-1)

    # 当输入真实标签，返回loss值；无真实标签，返回预测值
    def forward(self, x, mask=None, y=None):
        if y is not None:
            # 训练时，构建一个下三角的mask矩阵，让上下文之间没有交互
            print(mask.shape)
            x, _ = self.bert(x, attention_mask=mask)
            y_pred = self.classify(x)  # output shape:(batch_size, vocab_size)
            return self.loss(y_pred.view(-1, y_pred.shape[-1]), y.view(-1))
        else:
            # 预测时，可以不使用mask
            x, _ = self.bert(x)
            y_pred = self.classify(x)  # output shape:(batch_size, vocab_size)
            return torch.softmax(y_pred, dim=-1)


# 加载语料, 用title当成假想的prompt，content当成假想的answer
def load_corpus(path):
    corpus = []
    with open(path, encoding="utf8") as f:
        for line in f:
            line = json.loads(line)
            corpus.append([line["title"], line["content"]])
    return corpus

# 构造掩码，输入两个字符串的长度
def create_mask(s1, s2):
    len_s1 = s1 + 2  # cls + sep
    len_s2 = s2 + 1  # sep
    # 创建掩码张量
    mask = torch.ones(len_s1 + len_s2, len_s1 + len_s2)
    # 遍历s1的每个token
    for i in range(len_s1):
        # s1的当前token不能看到s2的任何token
        mask[i, len_s1:] = 0
        # 遍历s2的每个token
    for i in range(len_s2):
        # s2的当前token不能看到后面的s2 token
        mask[len_s1 + i, len_s1 + i + 1:] = 0
    return mask

def pad_mask(tensor, target_shape):
    # 获取输入张量和目标形状的长宽
    height, width = tensor.shape
    target_height, target_width = target_shape
    # 创建一个全零张量,形状为目标形状
    result = torch.zeros(target_shape, dtype=tensor.dtype, device=tensor.device)
    # 计算需要填充或截断的区域
    h_start = 0
    w_start = 0
    h_end = min(height, target_height)
    w_end = min(width, target_width)
    # 将原始张量对应的部分填充到全零张量中
    result[h_start:h_end, w_start:w_end] = tensor[:h_end - h_start, :w_end - w_start]
    return result

# sft的数据构造
# loss只计算答案部分，通过mask矩阵，让上下文之间没有交互
# label中使用-1，表示不参与训练
def build_dataset(tokenizer, corpus, max_length, batch_size):
    dataset = []
    for i, (prompt, answer) in enumerate(corpus):
        prompt_encode = tokenizer.encode(prompt, add_special_tokens=False)
        answer_encode = tokenizer.encode(answer, add_special_tokens=False)
        x = [tokenizer.cls_token_id] + prompt_encode + [tokenizer.sep_token_id] + answer_encode + [
            tokenizer.sep_token_id]
        y = len(prompt_encode) * [-1] + [-1] + answer_encode + [tokenizer.sep_token_id] + [-1]
        # 构建一个的mask矩阵，让prompt内可以交互，answer中上下文之间没有交互
        mask = create_mask(len(prompt_encode), len(answer_encode))
        # padding
        x = x[:max_length] + [0] * (max_length - len(x))
        y = y[:max_length] + [0] * (max_length - len(y))
        x = torch.LongTensor(x)
        y = torch.LongTensor(y)
        mask = pad_mask(mask, (max_length, max_length))
        dataset.append([x, mask, y])

    return DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=0)


# 建立模型
def build_model(vocab, char_dim, pretrain_model_path):
    model = LanguageModel(768, 21128, pretrain_model_path)
    return model

# 采样策略选择
def sampling_strategy(prob_distribution):
    if random.random() > 0.1:
        strategy = "greedy"
    else:
        strategy = "sampling"
    if strategy == "greedy":
        return int(torch.argmax(prob_distribution))
    elif strategy == "sampling":
        prob_distribution = prob_distribution.cpu().numpy()
        return np.random.choice(list(range(len(prob_distribution))), p=prob_distribution)

# 文本生成测试代码
def generate_sentence(openings, model, tokenizer):
    model.eval()
    openings = tokenizer.encode(openings)
    with torch.no_grad():
        # 生成文本超过30字则终止迭代
        while len(openings) <= 50:
            x = torch.LongTensor([openings])
            if torch.cuda.is_available():
                x = x.cuda()
            y = model(x)[0][-1]
            index = sampling_strategy(y)
            openings.append(index)
    return tokenizer.decode(openings)

def main(corpus_path, save_weight=True):
    epoch_num = 20  # 训练轮数
    batch_size = 32  # 每次训练样本个数
    char_dim = 768  # 每个字的维度
    max_length = 50  # 样本文本长度
    vocab_size = 21128  # 字表大小
    learning_rate = 0.001  # 学习率

    pretrain_model_path = r"F:\人工智能NLP\NLP资料\week6 语言模型\bert-base-chinese"
    tokenizer = BertTokenizer.from_pretrained(pretrain_model_path)

    corpus = load_corpus(corpus_path)  # 加载语料
    train_data = build_dataset(tokenizer, corpus, max_length, batch_size)  # 建立数据集
    model = build_model(vocab_size, char_dim, pretrain_model_path)  # 建立模型
    if torch.cuda.is_available():
        model = model.cuda()
    optim = torch.optim.Adam(model.parameters(), lr=learning_rate)  # 建立优化器
    print("文本词表模型加载完毕，开始训练")
    for epoch in range(epoch_num):
        model.train()
        watch_loss = []
        for x, mask, y in train_data:  # 构建一组训练样本
            if torch.cuda.is_available():
                x, mask, y = x.cuda(), mask.cuda(), y.cuda()
            optim.zero_grad()  # 梯度归零
            loss = model(x, mask, y)  # 计算loss
            loss.backward()  # 计算梯度
            optim.step()  # 更新权重
            watch_loss.append(loss.item())
        print("=========\n第%d轮平均loss:%f" % (epoch + 1, np.mean(watch_loss)))
        print(generate_sentence("北京明年拟推工作日半价观看电影", model, tokenizer))
        print(generate_sentence("南京一合金厂锅炉发生爆炸", model, tokenizer))
    if not save_weight:
        return
    else:
        base_name = os.path.basename(corpus_path).replace("txt", "pth")
        model_path = os.path.join("model", base_name)
        torch.save(model.state_dict(), model_path)
        return


if __name__ == "__main__":
    main("sample_data.json", False)