import math
from collections import defaultdict
import torch
from nltk.corpus import sentence_polarity
from torch import nn, optim
from torch.nn.utils.rnn import pad_sequence
from torch.utils.data import Dataset, DataLoader
from tqdm import tqdm
class Vocab:
def __init__(self,tokens = None):
# 使用列表存储所有的标记，从而可根据索引值获取相应的标记
self.idx_to_token = list()
# 使用字典实现标记到索引值的映射
self.token_to_idx = dict()
if tokens is not None:
if "" not in tokens:
tokens = tokens + [""]
for token in tokens:
self.idx_to_token.append(token)
self.token_to_idx[token] = len(self.idx_to_token) - 1
self.unk = self.token_to_idx[""]
# 创建词表，text包含若干句子，每个句子由若干标记构成
@classmethod
def build(cls,text,min_freq = 1 , reserved_tokens = None):
# 存储标记及其出现次数的映射字典
token_freqs = defaultdict(int)
# 无重复地进行标记
for sentence in text:
for token in sentence:
token_freqs[token] += 1
# 用户自定义的预留标记
uniq_tokens = [""] + (reserved_tokens if reserved_tokens else [])
uniq_tokens += [token for token,freq in token_freqs.items() if freq >= min_freq and token != ""]
return cls(uniq_tokens)
# 返回词表的大小，即词表中有多少个互不相同的标记
def __len__(self):
return len(self.idx_to_token)
# 查找输入标记对应的索引值
def __getitem__(self, token):
return self.token_to_idx.get(token,self.unk)
# 查找一系列输入标记对应的索引值
def convert_tokens_to_ids(self,tokens):
return [self[token] for token in tokens]
# 查找一系列索引值对应的标记
def convert_ids_to_tokens(self,indices):
return [self.idx_to_token[index] for index in indices]
# 创建存储数据的TransformerDataset的子类 → Transformer
class TransformerDataset(Dataset):
def __init__(self,data):
self.data = data
def __len__(self):
return len(self.data)
def __getitem__(self, i):
return self.data[i]
# 根据批次中每个序列长度生成Mask矩阵，以便处理长度不一致的序列，忽略掉比较短的序列的无效部分
def length_to_mask(lengths):
max_len = torch.max(lengths)
mask = torch.arange(max_len).expand(lengths.shape[0], max_len) < lengths.unsqueeze(1)
return mask
# 数据处理
def collate_fn(examples):
lengths = torch.tensor([len(ex[0]) for ex in examples])
inputs = [torch.tensor(ex[0]) for ex in examples]
targets = torch.tensor([ex[1] for ex in examples],dtype = torch.long)
# 对batch内的样本进行padding，使其具有相同长度
inputs = pad_sequence(inputs , batch_first= True)
return inputs , lengths , targets
# # 位置编码
class PositionalEncoding(nn.Module):
def __init__(self,d_model,dropout = 0.1,max_len = 512):
super(PositionalEncoding, self).__init__()
pe = torch.zeros(max_len,d_model)
position = torch.arange(0 , max_len , dtype= torch.float).unsqueeze(1)
div_term = torch.exp(torch.arange(0 , d_model , 2).float() * (-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).transpose(0 , 1)
self.register_buffer("pe",pe)
def forward(self,x):
x = x + self.pe[:x.size(0),:]
return x
# Transformer模型
class Transformer(nn.Module):
def __init__(self,vocab_size,embedding_dim,hidden_dim,num_class,dim_feedforward = 512,num_head = 2,num_layers = 2,dropout = 0.1,max_len = 128,activation:str = "relu"):
super(Transformer, self).__init__()
# 词嵌入层
self.embedding_dim = embedding_dim
self.embeddings = nn.Embedding(vocab_size,embedding_dim)
self.position_embedding = PositionalEncoding(embedding_dim,dropout,max_len)
# 编码层：Transformer
encoder_layer = nn.TransformerEncoderLayer(hidden_dim,num_head,dim_feedforward,dropout,activation)
self.transformer = nn.TransformerEncoder(encoder_layer,num_layers)
# 输出层
self.output = nn.Linear(hidden_dim,num_class)
def forward(self,inputs,lengths):
inputs = torch.transpose(inputs , 0 ,1)
hidden_states = self.embeddings(inputs)
hidden_states = self.position_embedding(hidden_states)
attention_mask = length_to_mask(lengths) == False
hidden_states = self.transformer(hidden_states,src_key_padding_mask = attention_mask)
hidden_states = hidden_states[0 , : , :]
output = self.output(hidden_states)
log_probs = torch.log_softmax(output , dim = 1)
return log_probs
# 设置超参数
embedding_dim = 128
hidden_dim = 128
num_class = 2
batch_size = 32
num_epoch = 5
# 加载数据
def load_sentence_polarity():
# 使用全部句子集合创建词表
vocab = Vocab.build(sentence_polarity.sents())
# 褒贬各4000句作为训练数据
train_data = [(vocab.convert_tokens_to_ids(sentence),0) for sentence in sentence_polarity.sents(categories="pos")[:4000]] + [
(vocab.convert_tokens_to_ids(sentence),1) for sentence in sentence_polarity.sents(categories = "neg")[:4000]
]
# 剩下数据作为测试数据
test_data = [(vocab.convert_tokens_to_ids(sentence), 0) for sentence in
sentence_polarity.sents(categories="pos")[4000:]] + [
(vocab.convert_tokens_to_ids(sentence), 1) for sentence in
sentence_polarity.sents(categories="neg")[4000:]
]
return train_data,test_data,vocab
train_data , test_data ,vocab = load_sentence_polarity()
train_dataset = TransformerDataset(train_data)
test_dataset = TransformerDataset(test_data)
train_data_loader = DataLoader(train_dataset,batch_size = batch_size,collate_fn = collate_fn,shuffle = True)
test_data_loader = DataLoader(test_dataset,batch_size = 1,collate_fn=collate_fn,shuffle = False)
# 加载模型
device = torch.device("cpu")
model = Transformer(len(vocab),embedding_dim,hidden_dim,num_class)
model.to(device)
# 训练过程
nll_loss = nn.NLLLoss()
optimizer = optim.Adam(model.parameters(),lr = 0.001)
model.train()
for epoch in range(num_epoch):
total_loss = 0
for batch in tqdm(train_data_loader,desc=f"Training Epoch{epoch}"):
inputs , lengths ,targets = [x.to(device) for x in batch]
log_probs = model(inputs,lengths)
loss = nll_loss(log_probs,targets)
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_loss += loss.item()
print(f"Loss:{total_loss:.2f}")
# 测试过程
acc = 0
for batch in tqdm(test_data_loader,desc = f"Testing"):
inputs , lengths , targets = [x.to(device) for x in batch]
with torch.no_grad():
output = model(inputs,lengths)
acc += (output.argmax(dim = 1) == targets).sum().item()
#输出在测试集上的准确率
print(f"Acc:{acc / len(test_data_loader):.2f}")
运行结果
