自然语言文本分类模型代码

以下是一个基于PyTorch的文本分类模型的示例代码,用于将给定的文本分为多个预定义类别:

python 复制代码
import torch
import torch.nn as nn
import torch.nn.functional as F

class TextClassifier(nn.Module):
    def __init__(self, vocab_size, embedding_dim, hidden_dim, output_dim, num_layers, bidirectional, dropout):
        super().__init__()
        
        self.embedding = nn.Embedding(vocab_size, embedding_dim)
        self.rnn = nn.LSTM(embedding_dim, hidden_dim, num_layers=num_layers, bidirectional=bidirectional, dropout=dropout)
        self.fc = nn.Linear(hidden_dim * 2 if bidirectional else hidden_dim, output_dim)
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, text, text_lengths):
        embedded = self.dropout(self.embedding(text))
        packed_embedded = nn.utils.rnn.pack_padded_sequence(embedded, text_lengths.to('cpu'), enforce_sorted=False)
        packed_output, (hidden, cell) = self.rnn(packed_embedded)
        output, output_lengths = nn.utils.rnn.pad_packed_sequence(packed_output)
        hidden = self.dropout(torch.cat((hidden[-2,:,:], hidden[-1,:,:]), dim=1) if self.rnn.bidirectional else hidden[-1,:,:])
        return self.fc(hidden.squeeze(0))

该模型将输入的文本作为整数序列传递给嵌入层,然后通过多层LSTM层进行处理,最终输出每个类别的预测概率。

在训练模型之前,需要将文本序列转换为整数标记,通常使用分词器/标记器完成此任务。另外还需要定义优化器和损失函数来训练模型。

以下是一个完整的训练脚本的示例:

python 复制代码
import torch.optim as optim
from torchtext.datasets import AG_NEWS
from torchtext.data.utils import get_tokenizer
from torchtext.vocab import build_vocab_from_iterator
from torch.utils.data import DataLoader
from torchtext.data.utils import ngrams_iterator
from torchtext.data.utils import get_tokenizer
from torch.utils.data.dataset import random_split
from collections import Counter

# 获取数据集和分词器
train_iter = AG_NEWS(split='train')
tokenizer = get_tokenizer('basic_english')

# 构建词汇表
counter = Counter()
for (label, line) in train_iter:
    counter.update(tokenizer(line))
vocab = build_vocab_from_iterator([counter])
vocab.set_default_index(vocab['<unk>'])

# 定义标记化函数和文本处理函数
def yield_tokens(data_iter):
    for _, text in data_iter:
        yield tokenizer(text)
        
def text_transform(tokenizer, vocab, data):
    """将文本数据转换为张量数据"""
    data = [vocab[token] for token in tokenizer(data)]
    return torch.tensor(data)
    
# 定义批次生成器
def collate_batch(batch):
    label_list, text_list, offsets = [], [], [0]
    for (_label, _text) in batch:
        label_list.append(_label-1)
        processed_text = torch.cat([text_transform(tokenizer, vocab, _text), torch.tensor([vocab['<eos>']])])
        text_list.append(processed_text)
        offsets.append(processed_text.size(0))
    label_list = torch.tensor(label_list)
    offsets = torch.tensor(offsets[:-1]).cumsum(dim=0)
    text_list = torch.cat(text_list)
    return label_list, text_list, offsets

# 构建数据集和数据加载器
train_iter, test_iter = AG_NEWS()
train_iter = list(train_iter)
test_iter = list(test_iter)
train_dataset = list(map(lambda x: (x[0], x[1]), train_iter))
test_dataset = list(map(lambda x: (x[0], x[1]), test_iter))
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True, collate_fn=collate_batch)
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=True, collate_fn=collate_batch)

# 创建模型和优化器
model = TextClassifier(len(vocab), 64, 128, 4, 2, True, 0.5)
optimizer = optim.Adam(model.parameters())

# 定义损失函数和训练函数
criterion = nn.CrossEntropyLoss()

def train(model, iterator, optimizer, criterion):
    epoch_loss = 0
    model.train()
    for (label, text, offsets) in iterator:
        optimizer.zero_grad()
        predictions = model(text, offsets)
        loss = criterion(predictions, label)
        loss.backward()
        optimizer.step()
        epoch_loss += loss.item()
    return epoch_loss / len(iterator)

# 训练模型
N_EPOCHS = 10
for epoch in range(N_EPOCHS):
    train_loss = train(model, train_loader, optimizer, criterion)
    print(f'Epoch: {epoch+1:02} | Train Loss: {train_loss:.3f}')

在训练过程结束后,可以使用该模型对新的文本进行分类。具体方法是将文本转换为整数标记序列,然后使用模型进行预测:

python 复制代码
# 对新文本进行分类
def predict(model, sentence):
    model.eval()
    tokenized = torch.tensor([vocab[token] for token in tokenizer(sentence)])
    length = torch.tensor([len(tokenized)])
    prediction = model(tokenized, length)
    return F.softmax(prediction, dim=1).detach().numpy()[0]

# 进行预测
test_sentence = "World markets are reacting to the news that the UK is set to leave the European Union."
pred_probs = predict(model, test_sentence)
print(pred_probs)

以上代码示例中使用了AG_NEWS数据集作为示例训练数据,可通过以下方式加载数据集:

python 复制代码
from torchtext.datasets import AG_NEWS
train_iter = AG_NEWS(split='train')
test_iter = AG_NEWS(split='test')

该数据集包含四个类别的新闻数据,每个类别各有120,000个训练示例和7,600个测试示例。完整的训练脚本和数据集可以在PyTorch官方文档中找到。

相关推荐
chaojihaomai3 小时前
AI电商商品图生成中的材质保真问题——扩散模型重建优先级分析与锁定式方案
人工智能·材质
大模型探索者5 小时前
2026 企业智能体平台选型指南:从通用到行业深耕,怎么选?
大数据·人工智能
深圳行云创新5 小时前
北极星知识图谱平台|太空越来越“拥挤”,知识图谱如何预测空间碎片碰撞风险?
人工智能·知识图谱
Am-Chestnuts5 小时前
国产AI多轮对话归档与多格式导出实践
人工智能·word
麻瓜生活睁不开眼6 小时前
Android16修改全局桌面视图边框四直角显示为弧边圆角
android·java·深度学习
workflower6 小时前
装备企业的 AI 路线图
人工智能·深度学习·机器学习·设计模式·机器人
AKAMAI6 小时前
Linode接口及默认防火墙现已全面开放使用
人工智能·云计算
前端开发江鸟7 小时前
RAG 回答错了,问题到底出在召回、重排,还是生成?
人工智能
美狐美颜SDK开放平台7 小时前
直播app开发如何实现主播级美颜效果?美颜sdk开发方案详解
人工智能·音视频·美颜sdk·视频美颜sdk·美颜api
东方小月7 小时前
从零开发一个 Coding Agent(四):使用状态机校验大模型事件流
前端·人工智能·后端