基于Mamba的拼音汉字转换模型

【图书推荐】《深入探索Mamba模型架构与应用》-优快云博客

深入探索Mamba模型架构与应用 - 商品搜索 - 京东

列到序列的转换是自然语言处理中的一个重要任务。在这个过程中，模型作为编码器，能够分别对序列中的个体以及整体进行编码。完成编码后，通过一个解码器，我们可以将这些编码信息转换成目标序列。这种转换不仅要求模型能够准确理解原始序列的语义，还需要它能够生成语法正确、语义连贯的目标序列。

为了实现这一目标，我们将深入探讨Mamba如何作为编码器和解码器进行工作，以及如何通过训练优化模型的性能。此外，我们还将讨论如何处理变长序列、如何引入注意力机制以提高转换的准确性，并探索一些先进的序列到序列的转换技术。

6.1.1  拼音汉字数据集详解与实战处理方法

首先进行数据集的准备和处理，在前面注意力章节的讲解中，我们已经遇到了拼音汉字数据集，本节将详细介绍这个数据集及其具体的处理方法。

1. 数据集展示

拼音汉字数据集如下：

A11_0	lv4 shi4 yang2 chun1 yan1 jing3 da4 kuai4 wen2 zhang1 de di3 se4 si4 yue4 de lin2 luan2 geng4 shi4 lv4 de2 xian1 huo2 xiu4 mei4 shi1 yi4 ang4 ran2	绿 是 阳 春 烟 景 大 块 文 章 的 底 色 四 月 的 林 峦 更 是 绿 得 鲜 活 秀 媚 诗 意 盎 然

A11_1	ta1 jin3 ping2 yao1 bu4 de li4 liang4 zai4 yong3 dao4 shang4 xia4 fan1 teng2 yong3 dong4 she2 xing2 zhuang4 ru2 hai3 tun2 yi1 zhi2 yi3 yi1 tou2 de you1 shi4 ling3 xian1	他 仅 凭 腰 部 的 力 量 在 泳 道 上 下 翻 腾 蛹 动 蛇 行 状 如 海 豚 一 直 以 一 头 的 优 势 领 先

A11_10	pao4 yan3 da3 hao3 le zha4 yao4 zen3 me zhuang1 yue4 zheng4 cai2 yao3 le yao3 ya2 shu1 de tuo1 qu4 yi1 fu2 guang1 bang3 zi chong1 jin4 le shui3 cuan4 dong4	炮 眼 打 好 了 炸 药 怎 么 装 岳 正 才 咬 了 咬 牙 倏 地 脱 去 衣 服 光 膀 子 冲 进 了 水 窜 洞

A11_100	ke3 shei2 zhi1 wen2 wan2 hou4 ta1 yi1 zhao4 jing4 zi zhi3 jian4 zuo3 xia4 yan3 jian3 de xian4 you4 cu1 you4 hei1 yu3 you4 ce4 ming2 xian3 bu4 dui4 cheng1	可 谁 知 纹 完 后 她 一 照 镜 子 只 见 左 下 眼 睑 的 线 又 粗 又 黑 与 右 侧 明 显 不 对 称

简单介绍一下。数据集中的数据分成3部分，每部分使用特定的空格键隔开：

A11_10 … … … ke3 shei2 … … …可 谁 … … …

1. 第一部分A11_i为序号，表示序列的条数和行号。
2. 其次是拼音编号，这里使用的是汉语拼音，与真实的拼音标注不同的是，去除了拼音原始标注，使用数字1、2、3、4替代，分别代表当前读音的第一声到第四声，这点请读者注意。
3. 最后一分部是汉字的序列，这里与第二部分的拼音部分一一对应。

2. 获取字库和训练数据

获取数据集中字库的个数也是一个非常重要的问题，一个非常好的办法是：使用set格式的数据读取全部字库中的不同字符。

创建字库和训练数据的完整代码如下：

max_length = 64
with open("zh.tsv", errors="ignore", encoding="UTF-8") as f:
    context = f.readlines()								#读取内容
    for line in context:		
        line = line.strip().split("	")					#切分每行中的不同部分
	   pinyin = ["GO"] + line[1].split(" ") + ["END"]	#处理拼音部分，在头尾加上起止符号
hanzi = ["GO"] + line[2].split(" ") + ["END"]	#处理汉字部分，在头尾加上起止符号
        for _pinyin, _hanzi in zip(pinyin, hanzi):	#创建字库
            pinyin_vocab.add(_pinyin)
hanzi_vocab.add(_hanzi)
pinyin = pinyin + ["PAD"] * (max_length - len(pinyin))
        hanzi = hanzi + ["PAD"] * (max_length - len(hanzi))
        pinyin_list.append(pinyin)				#创建拼音列表
hanzi_list.append(hanzi)						#创建汉字列表

这里做一个说明，首先context读取了全部数据集中的内容，之后根据空格将其分成3部分。对拼音和汉字部分，将其转换成一个序列，并在前后分别加上起止符GO和END。实际上也可以不加，为了明确地描述起止关系，从而加上了起止的标注。

实际上，还需要加上一个特定的符号PAD，这是为了对单行序列进行补全的操作，最终的数据如下：

['GO', 'liu2', 'yong3' , … … … , 'gan1', ' END', 'PAD', 'PAD' , … … …]
['GO', '柳', '永' , … … … , '感', ' END', 'PAD', 'PAD' , … … …]

pinyin_list和hanzi_list分别是两个列表，分别用来存放对应的拼音和汉字训练数据。最后不要忘记在字库中加上PAD符号。

pinyin_vocab = ["PAD"] + list(sorted(pinyin_vocab))
hanzi_vocab = ["PAD"] + list(sorted(hanzi_vocab))

3. 根据字库生成Token数据

获取的拼音标注和汉字标注的训练数据并不能直接用于模型训练，模型需要转换成Token的一系列数字列表，代码如下：

def get_dataset():
    pinyin_tokens_ids = []		# 新的拼音Token列表
    hanzi_tokens_ids = []		# 新的汉字Token列表

for pinyin,hanzi in zip(tqdm(pinyin_list),hanzi_list):
        #获取新的拼音Token
        pinyin_tokens_ids.append([pinyin_vocab.index(char) for char in pinyin])
        #获取新的汉字Token
        hanzi_tokens_ids.append([hanzi_vocab.index(char) for char in hanzi]) 

    return pinyin_vocab,hanzi_vocab,pinyin_tokens_ids,hanzi_tokens_ids

代码中创建了两个新的列表，分别对拼音和汉字的Token进行存储，而获取根据字库序号编号后新的序列Token。

6.1.2  Mamba模型的设计详解

在这里，我们使用Mamba架构的编码器完成模型的设计，在具体使用时，可以采用在第5章设计完成的Mamba模型进行模型的整理。代码如下：

import moudle

class Pin2Yin(torch.nn.Module):
    def _ _init_ _ (self, d_model=128, state_size=64, num_layers=3,
                 vocab_size=vocab_size, device=device):
        super()._ _init_ _()
        self.mamba_basic_model = moudle.Mamba(d_model=d_model, state_size=state_size, vocab_size=vocab_size,device=device, num_layers=num_layers)
        self.layer_norm = torch.nn.LayerNorm(d_model)
        self.logits_layer = torch.nn.Linear(d_model, vocab_size)

    def forward(self,x):
        embedding = self.mamba_basic_model(x)
        embedding = self.layer_norm(embedding)
        embedding = torch.nn.Dropout(0.1)(embedding)
        logits = self.logits_layer(embedding)
        return logits

在这里，我们直接使用moudle.Mamba作为编码器，对输入的信号进行编码。编码后的信号经过一个标准化层之后，再通过输出层将Embedding转换为输出向量。输出向量的维度要与字符数量相同，这是为了在softmax操作时，便于将向量转换为文字。

6.1.3  模型的训练与预测

最后，在使用我们设计的拼音汉字转换模型的基础上，完成最后的模型训练与预测任务，代码如下：

import torch.nn

import get_data

device = "cuda"
max_length = get_data.max_length
vocab_size = get_data.vocab_size

pinyin_tokens_ids,hanzi_tokens_ids = get_data.get_dataset()

import moudle

class Pin2Yin(torch.nn.Module):
    def _ _init_ _(self, d_model=128, state_size=64, num_layers=3,vocab_size=vocab_size, device=device):
        super()._ _init_ _()
        self.mamba_basic_model = moudle.Mamba(d_model=d_model, state_size=state_size, vocab_size=vocab_size,
                                              device=device, num_layers=num_layers)
        self.layer_norm = torch.nn.LayerNorm(d_model)
        self.logits_layer = torch.nn.Linear(d_model, vocab_size)


    def forward(self,x):
        embedding = self.mamba_basic_model(x)
        embedding = self.layer_norm(embedding)
        embedding = torch.nn.Dropout(0.1)(embedding)
        logits = self.logits_layer(embedding)

        return logits


if _ _name_ _ == '_ _main_ _':
    # 实例化模型，并移动到指定的设备上（CPU或GPU）
    model = Pin2Yin().to(device)

    # 定义优化器和损失函数
    optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5)
    loss_func = torch.nn.CrossEntropyLoss()

    #model_path = "./saver/modelpara.pt"
    #model.load_state_dict(torch.load(model_path), strict=True)

    # 设置训练参数
    batch_size = 384
    train_length = len(pinyin_tokens_ids)

    # 开始训练过程
    for epoch in range(6):
        train_num = train_length // batch_size
        train_loss, train_correct = 0, 0

        for i in range(train_num):
            start = i * batch_size
            end = (i + 1) * batch_size

            # 准备输入数据和标签
            batch_input_ids = torch.tensor(pinyin_tokens_ids[start:end],dtype=torch.long).to(device)
            batch_labels = torch.tensor(hanzi_tokens_ids[start:end],dtype=torch.long).to(device)

            # 进行前向传播和损失计算
            logits = model(batch_input_ids)

            loss = loss_func(logits.view(-1, logits.size(-1)), batch_labels.view(-1))

            # 进行反向传播和优化
            optimizer.zero_grad()
            loss.backward(retain_graph=True)
            optimizer.step()

            # 更新训练损失和准确率
            train_loss += loss.item()

        #计算并打印每个epoch的平均损失和准确率
        train_loss /= train_num
        print("epoch: ", epoch, "| train_loss:", train_loss)

        #if (epoch + 1) % 2 == 0:
        torch.save(model.state_dict(), "./saver/modelpara.pt")

    model.eval()
    vocab = get_data.vocab

    batch_input_ids = torch.tensor(pinyin_tokens_ids[19:29], dtype=torch.long).to(device)
    batch_labels = torch.tensor(hanzi_tokens_ids[19:29], dtype=torch.long).to(device)

    test_pred = model(torch.tensor(batch_input_ids))
    pred_label = torch.argmax(test_pred, dim=-1)
    for pred,true in zip(pred_label,batch_labels):
        pred = [vocab[idx] for idx in pred]
        true = [vocab[idx] for idx in true]
        print(pred)
        print(true)
        print("----------------------------")

结果如下：

epoch:  0 | train_loss: 0.34028792977333067
epoch:  1 | train_loss: 0.34051530361175536
epoch:  2 | train_loss: 0.34036597311496736
epoch:  3 | train_loss: 0.33997217416763303
epoch:  4 | train_loss: 0.3397156119346619
epoch:  5 | train_loss: 0.33966816365718844

['马', '晓', '年', '作', '画', '的', '特', '征', '与', '神', '韵', '其', '画', '风', '区', '别', '于', '文', '人', '画', '和', '年', '画', '可', '谓', '雅', '俗', '共', '赏', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD']
['马', '晓', '年', '作', '画', '的', '特', '征', '与', '神', '韵', '其', '画', '风', '区', '别', '于', '文', '人', '画', '和', '年', '画', '可', '谓', '雅', '俗', '共', '赏', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD']
...
['鉴', '于', '此', '经', '研', '究', '决', '定', '海', '峡', '情', '专', '栏', '原', '定', '三', '月', '底', '结', '束', '现', '延', '至', '五', '月', '三', '十', '一', '日', '特', '此', '告', '知', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD']
['鉴', '于', '此', '经', '研', '究', '决', '定', '海', '峡', '情', '专', '栏', '原', '定', '三', '月', '底', '结', '束', '现', '延', '至', '五', '月', '三', '十', '一', '日', '特', '此', '告', '知', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD', 'PAD']

Pin2Yin类继承自torch.nn.Module，并定义了网络结构。它包含一个名为mamba_basic_model的基础模型（可能是某种自定义的神经网络结构，来自moudle模块）、一个层归一化层（layer_norm）和一个线性层（logits_layer）用于输出预测。在forward方法中，输入数据首先通过基础模型进行编码，然后经过层归一化和Dropout层，最后通过线性层输出预测结果。

在_ _main_ _部分，代码首先实例化Pin2Yin模型，并将其移动到指定的设备上（GPU或CPU）。然后，定义优化器（AdamW）和损失函数（交叉熵损失）。在训练循环中，代码按照批次处理数据，进行前向传播、损失计算、反向传播和优化步骤。在每个epoch结束时，计算并打印平均训练损失，并保存模型参数。

最后，代码将模型设置为评估模式，并对一小批数据进行预测。预测结果和真实标签通过词汇表转换为汉字，并打印出来以供比较。整个代码流程覆盖了模型的定义、训练、保存和测试。