04

一、Transformer编码器项目

• 这里完成的是情感分类的一个项目，关于电影评价的

• 电商：

  • 客服的分类（售前、售后、商品质量、快递、投诉、价格）
  • 评论：正向的评论、负向的评论（你确实是给我五颗星，但是你说的话不太正向）

• 舆情监控：

  • 评论、视频、新闻稿子：是否有违法行为、政治问题

• 工程项目：代码不要求全部手写，合理借助 AI 开发，提高开发效率~但是一定要有自己的思想【设计项目、梳理流程、代码如何实现和解读】

• 后面可以用 BERT 网络来实现这个情感分析项目

1、数据准备和预处理

1.1 数据加载

• Internet Movie Database，IMDB 数据集是 50,000 条电影评论

• 数据集来自 Hugging Face：【https://huggingface.co/datasets/stanfordnlp/imdb 】

from datasets import load_dataset

dataset = load_dataset(
    "stanfordnlp/imdb",
    name="plain_text",
    cache_dir="./datasets/stanfordnlp/imdb",
    download_mode="reuse_dataset_if_exists",  # 缓存数据集
)
print(dataset)

1.2 数据清洗

• 根据自己的需求对数据进行初步的处理，如特殊字符替换等

import string

from datasets import load_dataset

dataset = load_dataset(
    "stanfordnlp/imdb",
    name="plain_text",
    cache_dir="./datasets/stanfordnlp/imdb",
    download_mode="reuse_dataset_if_exists",  # 缓存数据集
)
print(dataset)
# 训练集 'train'
train_data = dataset['train']
print(f"训练集样本数量: {len(train_data)}")
punctuation = string.punctuation
replacements = str.maketrans({p: "" for p in punctuation})
word_vocab = []
# 遍历选定的分片
for i in range(len(train_data)):
    # 访问方式改为 train_data[i]
    text = train_data[i]["text"].translate(replacements)
    word_vocab.extend(text.split())
print(word_vocab)

2、 词表构建

• 构建单词和 ID 之间的映射

import string
import pickle
from datasets import load_dataset

dataset = load_dataset(
    "stanfordnlp/imdb",
    name="plain_text",
    cache_dir="./datasets/stanfordnlp/imdb",
    download_mode="reuse_dataset_if_exists",
)

print(dataset)

train_data = dataset['train']
print(f"训练集样本数量: {len(train_data)}")

punctuation = string.punctuation
replacements = str.maketrans({p: "" for p in punctuation})

word_vocab = []

for i in range(len(train_data)):
    text = train_data[i]["text"].translate(replacements).lower()
    word_vocab.extend(text.split())

word_vocab = list(set(word_vocab))

word_vocab_dict = {word: i + 2 for i, word in enumerate(word_vocab)}
word_vocab_dict["<pad>"] = 0
word_vocab_dict["<unk>"] = 1

print("词表大小:", len(word_vocab_dict))

# ================= 保存 =================
with open("vocab.pkl", "wb") as f:
    pickle.dump(word_vocab_dict, f)

print("vocab 已保存为 vocab.pkl")

3、数据集加载器

import os
import string
from datasets import load_dataset
from torch.utils.data import Dataset

# 获取当前文件所在的目录
CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
# 通过string获取特殊符号
punctuation = string.punctuation
# 对特殊字符构造替换表
replacements = str.maketrans({p: "" for p in punctuation})

class IMDBDataset(Dataset):
    def __init__(self, split="train"):
        self.dataset = load_dataset(
            "stanfordnlp/imdb",
            name="plain_text",
            cache_dir=os.path.join(CURRENT_DIR, "datasets/stanfordnlp/imdb"),
            split=split,
            download_mode="reuse_dataset_if_exists",
        )

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

    def __getitem__(self, idx):
        text = self.dataset[idx]["text"].translate(replacements)
        label = self.dataset[idx]["label"]
        return text, label


if __name__ == '__main__':
    train_dataset = IMDBDataset(split="train")
    print(train_dataset)
    test_dataset = IMDBDataset(split="test")
    print(test_dataset)

4、数据整理

• 我们需要一个把一批样本打包成一个 batch 的函数，它负责对齐长度、处理 padding、转换为 tensor 等操作，使得 Transformer 模型可以正常训练

• 原因分析：文本数据是变长的，不像图像那样大小一致，语言数据句子长度一定是有差异的，比如：

["hello world"]
["hi"]
["this is a long sentence"]

• 把它们放进一个 batch 喂给模型时，显然需要对齐：

  • Tokenizer（分词/转 TokenID）

  • Pad 到相同长度

  • 构建 attention_mask，告诉模型哪些是 padding

• 默认的 $DataLoader$ 的 $collate_fn$ 是把一个 batch 的数据 $list$ 转成 $tensor$：

default_collate([x1, x2, x3]) => torch.tensor([x1, x2, x3])

• 但这在 NLP 中会直接报错，因为变长的句子转不了 tensor（维度不一致）
• 参考代码：

def collate_fn(batch):
    text_list, label_list = [], []

    for text, label in batch:
        ids = [word_vocab_dict.get(w, 1) for w in text.split()]

        # 截断
        ids = ids[:MAX_LEN]

        text_list.append(torch.tensor(ids))
        label_list.append(label)

    padded = pad_sequence(text_list, batch_first=True, padding_value=0)

    # mask：True表示padding
    mask = (padded == 0)

    return padded, mask, torch.tensor(label_list)

5、数据集加载器

• 在训练的时候，数据都是批次加载的

import os
import string
import torch
from datasets import load_dataset
from torch.nn.utils.rnn import pad_sequence
from torch.utils.data import Dataset, DataLoader

# 获取当前文件所在的目录
CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
# 通过string获取特殊符号
punctuation = string.punctuation
# 对特殊字符构造替换表
replacements = str.maketrans({p: "" for p in punctuation})
# 设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 统一长度限制
MAX_LEN = 500

class IMDBDataset(Dataset):
    def __init__(self, split="train"):
        self.dataset = load_dataset(
            "stanfordnlp/imdb",
            name="plain_text",
            cache_dir=os.path.join(CURRENT_DIR, "datasets/stanfordnlp/imdb"),
            split=split,
            download_mode="reuse_dataset_if_exists",  # 缓存数据集
        )

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

    def __getitem__(self, index):
        # 返回文本和标签
        return (
            self.dataset[index]["text"].translate(replacements),
            self.dataset[index]["label"],
        )


def main():
    # 训练集 'train'
    train_dataset = IMDBDataset(split="train")
    # 测试集 'test'
    test_dataset = IMDBDataset(split="test")
    # 遍历选定的分片
    word_vocab = []
    for i in range(len(train_dataset)):
        text, _ = train_dataset[i]
        word_vocab.extend(text.split())
    word_vocab = list(set(word_vocab))
    # 构建词表索引
    word_vocab_dict = {word: i + 1 for i, word in enumerate(word_vocab)}
    word_vocab_dict["<unk>"] = 0

    # 定义批处理函数
    def collate_fn(batch):
        text_list, label_list = [], []

        for text, label in batch:
            ids = [word_vocab_dict.get(w, 1) for w in text.split()]

            # 截断
            ids = ids[:MAX_LEN]

            text_list.append(torch.tensor(ids))
            label_list.append(label)

        padded = pad_sequence(text_list, batch_first=True, padding_value=0)

        # mask：True表示padding
        mask = (padded == 0)

        return padded, mask, torch.tensor(label_list)

    # 数据集加载器
    train_dataLoader = DataLoader(
        train_dataset,
        batch_size=4,
        shuffle=True,
        collate_fn=collate_fn,
    )
    test_dataLoader = DataLoader(
        test_dataset,
        batch_size=4,
        collate_fn=collate_fn,
    )
    for i, (text_batch, label_batch) in enumerate(train_dataLoader):
        print(text_batch.shape)
        if i == 2:
            break


if __name__ == '__main__':
    main()

6、位置编码

import math

import torch
from torch import nn



class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=500, dropout=0.1):
        super().__init__()
        self.dropout = nn.Dropout(dropout)

        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).unsqueeze(1)

        div_term = torch.exp(
            torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model)
        )

        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)

        pe = pe.unsqueeze(0)  # [1, max_len, d_model]
        self.register_buffer("pe", pe)

    def forward(self, x):
        # x: [batch, seq_len, dim]
        x = x + self.pe[:, : x.size(1), :]
        return self.dropout(x)

7、Transformer

from torch import nn
import torch
import math
from torch.nn import TransformerEncoderLayer, TransformerEncoder

MAX_LEN = 500  # 统一长度限制


class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=500, dropout=0.1):
        super().__init__()
        self.dropout = nn.Dropout(dropout)

        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).unsqueeze(1)

        div_term = torch.exp(
            torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model)
        )

        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)

        pe = pe.unsqueeze(0)  # [1, max_len, d_model]
        self.register_buffer("pe", pe)

    def forward(self, x):
        # x: [batch, seq_len, dim]
        x = x + self.pe[:, : x.size(1), :]
        return self.dropout(x)


class TransformerModel(nn.Module):
    def __init__(self, vocab_size, emb_dim, nhead, num_layers, num_class):
        super().__init__()

        self.embedding = nn.Embedding(vocab_size, emb_dim, padding_idx=0)
        self.pos_encoding = PositionalEncoding(emb_dim, MAX_LEN)

        encoder_layer = TransformerEncoderLayer(
            d_model=emb_dim,
            nhead=nhead,
            batch_first=True
        )

        self.encoder = TransformerEncoder(encoder_layer, num_layers=num_layers)

        self.fc = nn.Linear(emb_dim, num_class)

    def forward(self, x, mask):
        x = self.embedding(x)  # [B, L, D]
        x = self.pos_encoding(x)

        x = self.encoder(x, src_key_padding_mask=mask)

        x = x.mean(dim=1)

        return self.fc(x)

8、训练

import math
import os
import pickle
import string
import torch
from datasets import load_dataset
from torch import nn
from torch.nn import TransformerEncoderLayer, TransformerEncoder
from torch.nn.utils.rnn import pad_sequence
from torch.optim import lr_scheduler, Adam
from torch.utils.data import Dataset, DataLoader
from collections import Counter

# ================= 基础配置 =================
CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))

punctuation = string.punctuation
replacements = str.maketrans({p: "" for p in punctuation})

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

MAX_LEN = 500


# ================= Dataset =================
class IMDBDataset(Dataset):
    def __init__(self, split="train"):
        self.dataset = load_dataset(
            "stanfordnlp/imdb",
            name="plain_text",
            cache_dir=os.path.join(CURRENT_DIR, "datasets/stanfordnlp/imdb"),
            split=split,
            download_mode="reuse_dataset_if_exists",
        )

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

    def __getitem__(self, idx):
        text = self.dataset[idx]["text"].translate(replacements).lower()
        label = self.dataset[idx]["label"]
        return text, label


# ================= Positional Encoding =================
class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=500, dropout=0.1):
        super().__init__()
        self.dropout = nn.Dropout(dropout)

        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).unsqueeze(1)

        div_term = torch.exp(
            torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model)
        )

        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)

        self.register_buffer("pe", pe.unsqueeze(0))

    def forward(self, x):
        x = x + self.pe[:, :x.size(1), :]
        return self.dropout(x)


# ================= Transformer Model =================
class TransformerModel(nn.Module):
    def __init__(self, vocab_size, emb_dim, nhead, num_layers, num_class):
        super().__init__()

        self.embedding = nn.Embedding(vocab_size, emb_dim, padding_idx=0)
        self.pos_encoding = PositionalEncoding(emb_dim, MAX_LEN)

        encoder_layer = TransformerEncoderLayer(
            d_model=emb_dim,
            nhead=nhead,
            batch_first=True
        )

        self.encoder = TransformerEncoder(encoder_layer, num_layers=num_layers)

        self.fc = nn.Linear(emb_dim, num_class)

    def forward(self, x, mask):
        x = self.embedding(x)
        x = self.pos_encoding(x)

        x = self.encoder(x, src_key_padding_mask=mask)

        # ✔ 修复点：用 CLS-like pooling（比 mean 强很多）
        x = x[:, 0, :]

        return self.fc(x)


# ================= main =================
def main():

    # ================= 数据 =================
    train_dataset = IMDBDataset("train")
    test_dataset = IMDBDataset("test")

    # ================= 词表（修复核心问题） =================
    word_list = []

    for i in range(len(train_dataset)):
        text, _ = train_dataset[i]
        word_list.extend(text.split())

    # ✔ 用词频排序（稳定 + 工业标准）
    counter = Counter(word_list)
    vocab = [w for w, _ in counter.most_common(50000)]  # 限制 vocab

    word_vocab_dict = {w: i + 2 for i, w in enumerate(vocab)}
    word_vocab_dict["<pad>"] = 0
    word_vocab_dict["<unk>"] = 1

    vocab_size = len(word_vocab_dict)

    # ================= 保存词表 =================
    os.makedirs(os.path.join(CURRENT_DIR, "weights"), exist_ok=True)

    with open(os.path.join(CURRENT_DIR, "weights/vocab.pkl"), "wb") as f:
        pickle.dump(word_vocab_dict, f)

    print("vocab size:", vocab_size)

    # ================= collate =================
    def collate_fn(batch):
        text_list, label_list = [], []

        for text, label in batch:
            tokens = text.split()
            ids = [word_vocab_dict.get(w, 1) for w in tokens][:MAX_LEN]

            text_list.append(torch.tensor(ids))
            label_list.append(label)

        padded = pad_sequence(text_list, batch_first=True, padding_value=0)
        mask = (padded == 0)

        return padded, mask, torch.tensor(label_list)

    # ================= DataLoader =================
    train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, collate_fn=collate_fn)
    test_loader = DataLoader(test_dataset, batch_size=32, collate_fn=collate_fn)

    # ================= model =================
    model = TransformerModel(
        vocab_size=vocab_size,
        emb_dim=256,
        nhead=2,
        num_layers=2,
        num_class=2
    ).to(device)

    criterion = nn.CrossEntropyLoss()
    optimizer = Adam(model.parameters(), lr=1e-4)
    scheduler = lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.5)

    # ================= train =================
    def train():
        model.train()
        total_loss, total_acc, total_n = 0, 0, 0

        for x, mask, y in train_loader:
            x, mask, y = x.to(device), mask.to(device), y.to(device)

            pred = model(x, mask)
            loss = criterion(pred, y)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            total_loss += loss.item() * y.size(0)
            total_acc += (pred.argmax(1) == y).sum().item()
            total_n += y.size(0)

        return total_loss / total_n, total_acc / total_n

    # ================= eval =================
    def evaluate():
        model.eval()
        total_loss, total_acc, total_n = 0, 0, 0

        with torch.no_grad():
            for x, mask, y in test_loader:
                x, mask, y = x.to(device), mask.to(device), y.to(device)

                pred = model(x, mask)
                loss = criterion(pred, y)

                total_loss += loss.item() * y.size(0)
                total_acc += (pred.argmax(1) == y).sum().item()
                total_n += y.size(0)

        return total_loss / total_n, total_acc / total_n

    # ================= train loop =================
    best_acc = 0

    for epoch in range(100):

        train_loss, train_acc = train()
        test_loss, test_acc = evaluate()

        scheduler.step()

        print(f"Epoch {epoch+1}: "
              f"train_acc={train_acc:.4f}, test_acc={test_acc:.4f}")

        # save best
        if test_acc > best_acc:
            best_acc = test_acc
            torch.save(model.state_dict(), os.path.join(CURRENT_DIR, "weights/best_model.pt"))
            print("✅ saved best model")

    print("训练完成")


if __name__ == "__main__":
    main()

• 训练过程：

Epoch 1: train_acc=0.6370, test_acc=0.7514
✅ saved best model
Epoch 2: train_acc=0.7801, test_acc=0.8056
✅ saved best model
Epoch 3: train_acc=0.8243, test_acc=0.8249
✅ saved best model
Epoch 4: train_acc=0.8433, test_acc=0.8323
✅ saved best model
Epoch 5: train_acc=0.8639, test_acc=0.8380
✅ saved best model
Epoch 6: train_acc=0.8877, test_acc=0.8429
✅ saved best model
Epoch 7: train_acc=0.8989, test_acc=0.8460
✅ saved best model
Epoch 8: train_acc=0.9052, test_acc=0.8496
✅ saved best model
Epoch 9: train_acc=0.9146, test_acc=0.8479
Epoch 10: train_acc=0.9209, test_acc=0.8485
Epoch 11: train_acc=0.9305, test_acc=0.8510
✅ saved best model
Epoch 12: train_acc=0.9330, test_acc=0.8470
Epoch 13: train_acc=0.9384, test_acc=0.8482
Epoch 14: train_acc=0.9412, test_acc=0.8455
Epoch 15: train_acc=0.9424, test_acc=0.8430
Epoch 16: train_acc=0.9477, test_acc=0.8504
Epoch 17: train_acc=0.9518, test_acc=0.8492
.....

二、推理

• 加载训练的模型，测试推理【训练轮次不是很多，可能效果不好】
• 参考代码

import pickle
import torch
from torch import nn
import numpy as np
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
MAX_LEN = 500
# 关闭科学计数法
np.set_printoptions(suppress=True)


class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=500, dropout=0.1):
        super().__init__()
        self.dropout = nn.Dropout(dropout)

        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).unsqueeze(1)

        div_term = torch.exp(
            torch.arange(0, d_model, 2) * (-torch.log(torch.tensor(10000.0)) / d_model)
        )

        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)

        self.register_buffer("pe", pe.unsqueeze(0))

    def forward(self, x):
        x = x + self.pe[:, :x.size(1), :]
        return self.dropout(x)


class TransformerModel(nn.Module):
    def __init__(self, vocab_size, emb_dim, nhead, num_layers, num_class):
        super().__init__()

        self.embedding = nn.Embedding(vocab_size, emb_dim, padding_idx=0)
        self.pos_encoding = PositionalEncoding(emb_dim, MAX_LEN)

        encoder_layer = nn.TransformerEncoderLayer(
            d_model=emb_dim,
            nhead=nhead,
            batch_first=True
        )

        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        self.fc = nn.Linear(emb_dim, num_class)

    def forward(self, x, mask):
        x = self.embedding(x)
        x = self.pos_encoding(x)

        x = self.encoder(x, src_key_padding_mask=mask)
        x = x.mean(dim=1)

        return self.fc(x)


def predict(text):
    # ===== 加载词表 =====
    with open(r"F:\workspace\AI\stu_nlp\transfomrer情感分析\weights\vocab.pkl", "rb") as f:
        word_vocab_dict = pickle.load(f)

    vocab_size = len(word_vocab_dict)

    # ===== 初始化模型 =====
    model = TransformerModel(
        vocab_size=vocab_size,
        emb_dim=256,
        nhead=2,
        num_layers=2,
        num_class=2
    ).to(device)

    # ===== 加载权重 =====
    model.load_state_dict(
        torch.load(r"F:\workspace\AI\stu_nlp\transfomrer情感分析\weights\best_model.pt", map_location=device, weights_only=False))
    model.eval()
    # 文本预处理
    tokens = text.split()
    ids = [word_vocab_dict.get(w, 1) for w in tokens]
    ids = ids[:MAX_LEN]

    x = torch.tensor(ids).unsqueeze(0).to(device)
    mask = (x == 0)
    with torch.no_grad():
        logits = model(x, mask)
        prob = torch.softmax(logits, dim=1)
        pred = torch.argmax(prob, dim=1).item()
    label = "正面 😊" if pred == 1 else "负面 😠"
    return label, prob.cpu().numpy()


if __name__ == '__main__':
    print(predict("this movie is very good"))
    print(predict("this movie is very bad"))

• 运行结果

('正面 😊', array([[0.00198748, 0.99801254]], dtype=float32))
('负面 😠', array([[0.96143734, 0.03856267]], dtype=float32))