【AI-模型代码解析】中文问题生成模型

2024-10-20

字数统计: 2.5k | 阅读时长≈ 14 分钟

¶写在前面

¶代码总览

_train_0.py

## 自动写诗的例子
import io
import sys
sys.path.append("/NLP/bert_seq2seq-master/RoBERTa")
import torch
from tqdm import tqdm
import numpy as np
import json
from config import sentiment_batch_size, sentiment_lr, roberta_chinese_model_path
from model.test_n_0 import Seq2SeqModel
from model.roberta_model_0 import BertConfig
import time
from torch.utils.data import Dataset, DataLoader
from tokenizer import Tokenizer, load_chinese_base_vocab, BasicTokenizer
from lattice.utils_ import Trie, get_skip_path
from Data.lattice.TestLattice import load_cival_rules_rich_pretrain_word_list

w_trie = Trie()
def get_w_tire():
    w_list = load_cival_rules_rich_pretrain_word_list("./Data/lattice/toumu.txt",
                                                      _refresh=False,
                                                      _cache_fp='cache/{}'.format("rules_lattice")
                                                      )
    for w in w_list:  # 构建词典树
        w_trie.insert(w)
    print(w_trie)
# get_w_tire()
def create_dataset():
    get_w_tire()
    input_result = []
    output_result = []
    token_len = []
    target = io.open("./Data/new_tumu/question_train", encoding='UTF-8')
    source = io.open("./Data/new_tumu/answer_train", encoding='UTF-8')
    for scr, tar in zip(source, target):
        scr = scr.replace("\n", "")
        tar = tar.replace("\n", "")

        if len(scr) + len(tar) >= 100:
            continue
        lexicons = get_skip_path(scr, w_trie)
        for i in range(len(lexicons) - 1, 0, -1):
            for j in range(i - 1, -1, -1):
                if lexicons[i][2] == lexicons[j][2]:
                    del lexicons[i]
                    break
        # if len(lexicons) > 21:
        #     lexicons = lexicons[0:20]
        tempLexicons = list(map(lambda x: x[2], lexicons))
        lexicons.append(tempLexicons)
        lexicons.insert(0, scr)
        input_result.append(lexicons)

        output_result.append(tar)
    print(len(input_result))
    return input_result, output_result

## 自定义dataset
class PoemDataset(Dataset):
    """
    针对特定数据集，定义一个相关的取数据的方式
    """

    def __init__(self):
        ## 一般init函数是加载所有数据
        super(PoemDataset, self).__init__()
        # 读原始数据
        self.sents_src, self.sents_tgt = create_dataset()
        self.word2idx = load_chinese_base_vocab()
        self.idx2word = {k: v for v, k in self.word2idx.items()}
        self.tokenizer = Tokenizer(self.word2idx)
        self.pidding_idx = 0
        # print(self.sents_src[:3])

    def deal(self, i):
        src = []
        l = len(self.sents_src[i][-1])
        len_position = self.tokenizer._tokenize(self.sents_src[i][0])

        relation_position_s = [self.pidding_idx]+list(range(1, len(len_position)+1))
        relation_position_e = [self.pidding_idx]+list(range(1, len(len_position)+1))
        pre_start = 0
        start = 0
        for index in range(len(self.sents_src[i]) - 1):
            temp = self.sents_src[i][index]
            if index == 0:
                src.append(temp)
                continue
            isAppend = 0
            for k in temp[2]:
                flat = len(temp[2])
                if k == temp[2][0]:
                    try:
                        if pre_start == 0 or pre_start < 0:
                            start = len_position.index(k)+1
                        else:
                            start = len_position.index(k, pre_start)+1
                    except:
                        break
                    relation_position_s.append(start)
                    relation_position_e.append(self.pidding_idx)
                if k == temp[2][-1]:
                    try:
                        if pre_start == 0 or pre_start < 0:
                            end = len_position.index(k)+1
                        else:
                            end = len_position.index(k, pre_start)+1
                            if end - start != flat-1:
                                isAppend = 1
                                for m in range(len(temp[2]) - 1):
                                    relation_position_s.pop(-1)
                                    relation_position_e.pop(-1)
                                break
                        pre_start = end - 2
                    except:
                        for m in range(len(temp[2])-1):
                            relation_position_s.pop(-1)
                            relation_position_e.pop(-1)
                        break
                    relation_position_e.append(end)
                    relation_position_s.append(self.pidding_idx)
                elif k != temp[2][0] and k != temp[2][-1]:
                    relation_position_s.append(self.pidding_idx)
                    relation_position_e.append(self.pidding_idx)
            if isAppend == 0:
                src.append(temp[2])
        src = ''.join(src)

        tgt = self.sents_tgt[i]
        # src = self.sents_src[i][0]


        return src, tgt, relation_position_s, relation_position_e, l

    def __getitem__(self, i):
        ## 得到单个数据
        src, tgt, relation_position_s, relation_position_e, lex_num = self.deal(i)
        token_ids, token_type_ids = self.tokenizer.encode(src, tgt)
        # print(token_type_ids)
        input_length = token_type_ids[:token_type_ids.index(1)]
        # if len(token_ids) < len(src):
        #     relation_position_s = list(range(len(token_ids)))
        #     relation_position_e = list(range(len(token_ids)))

        output = {
            "token_ids": token_ids,
            "token_type_ids": token_type_ids,
            "relation_position_s": relation_position_s,
            "relation_position_e": relation_position_e,
            "lex_num": lex_num,
            "input_length": input_length
            }
        return output

    def __len__(self):

        return len(self.sents_src)



# class PoemDataset2(Dataset):
#     """
#     针对特定数据集，定义一个相关的取数据的方式
#     """
#
#     def __init__(self):
#         ## 一般init函数是加载所有数据
#         super(PoemDataset2, self).__init__()
#         # 读原始数据
#         self.sents_src, self.sents_tgt = create_dataset2()
#         self.word2idx = load_chinese_base_vocab()
#         self.idx2word = {k: v for v, k in self.word2idx.items()}
#         self.tokenizer = Tokenizer(self.word2idx)
#         # print(self.sents_src[:3])
#
#     def deal(self, i):
#         src = []
#         l = len(self.sents_src[i][-1])
#         relation_position_s = list(range(len(self.sents_src[i][0])))
#         relation_position_e = list(range(len(self.sents_src[i][0])))
#         for index in range(len(self.sents_src[i]) - 1):
#             temp = self.sents_src[i][index]
#             if index == 0:
#                 src.append(temp)
#                 continue
#             src.append(temp[2])
#             relation_position_s.append(temp[0])
#             relation_position_e.append(temp[1])
#         src = ''.join(src)
#         tgt = self.sents_tgt[i]
#
#         return src, tgt, relation_position_s, relation_position_e, l
#
#     def __getitem__(self, i):
#         ## 得到单个数据
#         # print(i)
#         # print(self.sents_src[i][0])
#         # print(self.sents_src[i])
#         # print(len(self.sents_src[i][0]))
#         # print(len(self.sents_src[i][-1]))
#         src, tgt, relation_position_s, relation_position_e, lex_num = self.deal(i)
#         token_ids, token_type_ids = self.tokenizer.encode(src, tgt)
#
#         if len(token_ids) < len(src):
#             relation_position_s = list(range(len(token_ids)))
#             relation_position_e = list(range(len(token_ids)))
#
#         output = {
#             "token_ids": token_ids,
#             "token_type_ids": token_type_ids,
#             "relation_position_s": relation_position_s,
#             "relation_position_e": relation_position_e,
#             "lex_num": lex_num
#         }
#         return output
#
#     def __len__(self):
#
#         return len(self.sents_src)
class PoemDataset1(Dataset):
    """
    针对特定数据集，定义一个相关的取数据的方式
    """

    def __init__(self):
        ## 一般init函数是加载所有数据
        super(PoemDataset1, self).__init__()
        # 读原始数据
        self.sents_src, self.sents_tgt = create_dataset()
        self.word2idx = load_chinese_base_vocab()
        self.idx2word = {k: v for v, k in self.word2idx.items()}
        self.tokenizer = Tokenizer(self.word2idx)

    def deal(self, i):
        src = []
        l = len(self.sents_src[i][-1])
        relation_position_s = list(range(len(self.sents_src[i][0])))
        relation_position_e = list(range(len(self.sents_src[i][0])))
        for index in range(len(self.sents_src[i]) - 1):
            temp = self.sents_src[i][index]
            if index == 0:
                src.append(temp)
                continue
            src.append(temp[2])
            relation_position_s.append(temp[0])
            relation_position_e.append(temp[1])
        src = ''.join(src)
        tgt = self.sents_tgt[i]

        return src, tgt, relation_position_s, relation_position_e, l

    def __getitem__(self, i):
        ## 得到单个数据
        # print(i)
        src, tgt, relation_position_s, relation_position_e, lex_num = self.deal(i)
        token_ids, token_type_ids = self.tokenizer.encode(src, tgt)

        if len(token_ids) < len(src):
            relation_position_s = list(range(len(token_ids)))
            relation_position_e = list(range(len(token_ids)))

        output = {
            "token_ids": token_ids,
            "token_type_ids": token_type_ids,
            "relation_position_s": relation_position_s,
            "relation_position_e": relation_position_e,
            "lex_num": lex_num
        }
        return output

    def __len__(self):

        return len(self.sents_src)


def collate_fn(batch):
    """
    动态padding， batch为一部分sample
    """

    def padding(indice, max_length, pad_idx=0):
        """
        pad 函数
        注意 token type id 右侧pad是添加1而不是0，1表示属于句子B
        """
        pad_indice = [item + [pad_idx] * max(0, max_length - len(item)) for item in indice]
        return torch.tensor(pad_indice)

    token_ids = [data["token_ids"] for data in batch]
    max_length = max([len(t) for t in token_ids])
    token_type_ids = [data["token_type_ids"] for data in batch]
    relation_position_s = [data["relation_position_s"] for data in batch]
    relation_position_e = [data["relation_position_e"] for data in batch]
    lex_num = [data["lex_num"] for data in batch]
    input_length = [data["input_length"] for data in batch]
    padding_idx = 0

    token_ids_padded = padding(token_ids, max_length)
    token_type_ids_padded = padding(token_type_ids, max_length, pad_idx=1)
    relation_position_s_padded = padding(relation_position_s, max_length, pad_idx=padding_idx)
    relation_position_e_padded = padding(relation_position_e, max_length, pad_idx=padding_idx)
    target_ids_padded = token_ids_padded[:, 1:].contiguous()


    return token_ids_padded, token_type_ids_padded, target_ids_padded, relation_position_s_padded, relation_position_e_padded, lex_num, input_length
    # return token_ids_padded, token_type_ids_padded, target_ids_padded, relation_position_s_padded, relation_position_e_padded, lex_num


class PoemTrainer:
    def __init__(self, random):
        # 加载情感分析数据
        self.random = random
        self.pretrain_model_path = roberta_chinese_model_path
        # 这个最近模型的路径可以用来继续训练，而不是每次从头训练
        # self.recent_model_path = "../poem_state_dict/bert_poem.model.epoch.9"
        self.batch_size = sentiment_batch_size
        self.lr = sentiment_lr
        # 加载字典
        self.word2idx = load_chinese_base_vocab()
        # 判断是否有可用GPU
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        print("device: " + str(self.device))
        # 定义模型超参数
        bertconfig = BertConfig(vocab_size=len(self.word2idx))
        # 初始化BERT模型
        self.bert_model = Seq2SeqModel(config=bertconfig)
        ## 加载预训练的模型～
        self.load_model(self.bert_model, self.pretrain_model_path)
        # self.load_recent_model(self.bert_model, self.recent_model_path)
        # 将模型发送到计算设备(GPU或CPU)
        self.bert_model.to(self.device)
        # 声明需要优化的参数
        self.optim_parameters = list(self.bert_model.parameters())
        # self.init_optimizer(lr=self.lr)
        self.init_optimizer(self.lr)
        # self.freeze_parameters()
        # 声明自定义的数据加载器
        dataset = PoemDataset()
        # dataset1 = PoemDataset1()
        self.dataloader = DataLoader(dataset, batch_size=self.batch_size, shuffle=True, collate_fn=collate_fn)
        # self.dataloader1 = DataLoader(dataset1, batch_size=self.batch_size, shuffle=True, collate_fn=collate_fn)

    def init_optimizer(self, lr):
        # 用指定的学习率初始化优化器
        # 冻结某层参数
        for name, param in self.bert_model.named_parameters():
            param.requires_grad = True
        print("模型参数解冻")
        self.optimizer = torch.optim.Adam(self.optim_parameters, lr=lr, weight_decay=1e-3)

    # def init_optimizer(self, lr):
    #     # 用指定的学习率初始化优化器
    #     self.optimizer = torch.optim.Adam(self.optim_parameters, lr=lr, weight_decay=1e-3)

    # def load_model(self, model, pretrain_model_path):

    def load_model(self, model, pretrain_model_path):

        checkpoint = torch.load(pretrain_model_path)
        # 模型刚开始训练的时候, 需要载入预训练的BERT

       # checkpoint = {k[5:]: v for k, v in checkpoint.items()
       #               if k[:4] == "bert" and "pooler" not in k}

        model.load_state_dict(checkpoint, strict=False)
        torch.cuda.empty_cache()
        print("{} loaded!".format(pretrain_model_path))

    def load_recent_model(self, model, recent_model_path):
        checkpoint = torch.load(recent_model_path)
        model.load_state_dict(checkpoint)
        torch.cuda.empty_cache()
        print(str(recent_model_path) + "loaded!")

    def train(self, epoch):
        # 一个epoch的训练
        self.bert_model.train()

        self.iteration(epoch, train=True)

    def freeze_parameters(self):
        freeze_layers = ['embeddings', 'layer.0', 'layer.1', 'layer.2', 'layer.3', 'layer.4', 'layer.5', 'layer.6', 'layer.7', 'layer.8', 'layer.9', 'layer.10', 'layer.11', 'bert.pooler', 'out.weight', 'out.bias']

        for name, param in self.bert_model.named_parameters():
            param.requires_grad = True
            for ele in freeze_layers:
                if ele in name:
                    param.requires_grad = False
                    break
        print("冻结模型参数")
        self.optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, self.bert_model.parameters()), lr=self.lr, weight_decay=1e-3)

    def iteration(self, epoch, train=True):
        total_loss = 0
        total_loss1 = 0
        start_time = time.time()  ## 得到当前时间
        step = 0
        # 训练开放域

        # for token_ids, token_type_ids, target_ids, relation_position_s, relation_position_e, lex_num in tqdm(self.dataloader, position=0, leave=True):
        for token_ids, token_type_ids, target_ids, relation_position_s, relation_position_e, lex_num ,input_length in tqdm(self.dataloader, position=0, leave=True):
            step += 1
            if step % 800== 0:
                self.bert_model.eval()
                test_data = ["设备限界车辆在故障运行状态下所形成的最大动态包络线，用以限制行车区的设备安装。"]
                tail = -1
                for text in test_data:
                    print(self.bert_model.generate(text, beam_size=1, tail=tail, device=self.device))
                    # tail = 0
                self.bert_model.train()

            token_ids = token_ids.to(self.device)

            token_type_ids = token_type_ids.to(self.device)
            target_ids = target_ids.to(self.device)

            relation_position_s = relation_position_s.to(self.device)
            relation_position_e = relation_position_e.to(self.device)


            # 因为传入了target标签，因此会计算loss并且返回
            # try:
            predictions, loss = self.bert_model(token_ids,
                                                token_type_ids,
                                                labels=target_ids,
                                                device=self.device,
                                                random=-1,
                                                lex_num=lex_num,
                                                pos_s=relation_position_s,
                                                pos_e=relation_position_e,
                                                input_length=input_length
                                                )
            if train:
                # 清空之前的梯度
                self.optimizer.zero_grad()
                # 反向传播, 获取新的梯度
                loss.backward()
                # 用获取的梯度更新模型参数
                self.optimizer.step()

            # 为计算当前epoch的平均loss
            total_loss += loss.item()

        with open("losses", "a") as f:
            epoch_loss = np.mean(total_loss)
            # epoch_loss1 = np.mean(total_loss1)
            f.write("Epoch: {}, Loss_1t: {}\n".format(str(epoch), str(epoch_loss)))
        end_time = time.time()
        spend_time = end_time - start_time

        # 打印训练信息
        # print("epoch is " + str(epoch) + ". loss_1 is " + str(total_loss) + ". loss_2 is " + str(total_loss1) + ". spend time is " + str(spend_time))
        print("epoch is " + str(epoch) + ". loss_1 is " + str(total_loss) + ". spend time is " + str(spend_time))
        # 保存模型
        # self.bert_model.eval()
        # test_data = ["适用于地基差的场地，但耐腐蚀性差，需经常维护。"]
        # for text in test_data:
        #     print(self.bert_model.generate(text, beam_size=2, device=self.device))
        # self.bert_model.train()
        self.save_state_dict(self.bert_model, epoch)
        # if epoch%10 == 0:
        #     epoch_0 = epoch % 5
        #     checkpoint = torch.load("./test_1/bert_poem.model.epoch-0.{}".format(str(epoch_0)))
        #
        #     for k, v in checkpoint.items():
        #         if "Transformer_Encoder.layer_0.pe_ss" in k:
        #             print(v)

    def save_state_dict(self, model, epoch, file_path="bert_poem.model"):
        """存储当前模型参数"""
        epoch = epoch % 5
        save_path = "./new12-2/" + file_path + ".epoch-0.{}".format(str(epoch))
        torch.save(model.state_dict(), save_path)
        print("{} saved!".format(save_path))


if __name__ == '__main__':

    # word2idx = load_chinese_base_vocab()
    # tokenier = Tokenizer(word2idx)
    #
    trainer = PoemTrainer(True)
    train_epoches = 500
    for epoch in range(1, train_epoches + 1):
        # 训练一个epoch
        torch.cuda.empty_cache()
        trainer.train(epoch)

该代码实现了基于BERT的自动写诗模型
主要流程包括：
数据加载、模型定义、训练和测试

¶代码分部分解析

¶导入必要的库和模块

import io
import sys
sys.path.append("/NLP/bert_seq2seq-master/RoBERTa")
import torch
from tqdm import tqdm
import numpy as np
import json
from config import sentiment_batch_size, sentiment_lr, roberta_chinese_model_path
from model.test_n_0 import Seq2SeqModel
from model.roberta_model_0 import BertConfig
import time
from torch.utils.data import Dataset, DataLoader
from tokenizer import Tokenizer, load_chinese_base_vocab, BasicTokenizer
from lattice.utils_ import Trie, get_skip_path
from Data.lattice.TestLattice import load_cival_rules_rich_pretrain_word_list

¶构建词汇树

w_trie = Trie()

def get_w_tire():
    w_list = load_cival_rules_rich_pretrain_word_list("./Data/lattice/toumu.txt",
                                                      _refresh=False,
                                                      _cache_fp='cache/{}'.format("rules_lattice")
                                                      )
    for w in w_list:  # 构建词典树
        w_trie.insert(w)
    print(w_trie)

这部分代码创建了一个词汇树（Trie）
并从指定文件中加载预训练词汇列表来填充这个树
这有助于后续处理文本时进行词汇匹配