NLP-文本分类-MLP
import os
import torch
import argparse
import torch.nn as nn
from pathlib import Path
from sklearn.metrics import accuracy_score
from torch.utils.data import Dataset, DataLoader
from torch.optim import Adam
import pickle as pkl
import warnings
warnings.filterwarnings('ignore')
path = Path(__file__).parent
def read_data(data_str):
text, label, max_len = [], [], []
with open(os.path.join(args.data_dif, f'{data_str}.txt'), 'r', encoding='utf-8') as f:
for line in f.readlines():
if not line:
continue
text_i, label_i = line.strip().split('\t')
text.append(text_i)
label.append(label_i)
max_len.append(len(text_i))
return text, label, max(max_len)
class Datasets(Dataset):
def __init__(self, text, label, word2idx, len_max):
self.text = text
self.label = label
self.word2idx = word2idx
self.len_max = len_max
self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
def __len__(self):
return len(self.text)
def __getitem__(self, idx):
text_ = self.text[idx][:self.len_max]
text_id = [self.word2idx.get(text_i, 1) for text_i in text_]
text_id = text_id + [0] * (self.len_max - len(text_id))
text_id = torch.tensor(text_id).to(self.device)
label_ = self.label[idx]
return text_id, torch.tensor(int(label_)).to(self.device)
def build_dict(train_text, embedding_size):
word2idx = {"<PAD>": 0, "<UNK>": 1}
for text_i in train_text:
for word in text_i:
word2idx[word] = word2idx.get(word, len(word2idx))
embedding = nn.Embedding(len(word2idx), embedding_size)
pkl.dump([word2idx,embedding], open(args.emb_words, 'wb'))
return word2idx, embedding
class Mymodel(nn.Module):
def __init__(self, vocal_size, embedding_matrix, hidden_size, class_num, dropout_rage=0.1):
super(Mymodel, self).__init__()
self.vocal_size = vocal_size
self.embedding_size = embedding_matrix.weight.shape[1]
self.hidden_size = hidden_size
self.class_num = class_num
self.emb = embedding_matrix
self.linear = nn.Linear(self.embedding_size, class_num)
self.doupout = nn.Dropout(dropout_rage)
def forward(self, x):
x = self.emb(x)
x = x.mean(dim=1)
out = self.linear(x)
return out
def generate(args,texts):
device = 'cuda' if torch.cuda.is_available() else 'cpu'
train_text, train_label, len_max = read_data('train')
if os.path.exists(args.emb_words):
datas = pkl.load(open(args.emb_words, 'rb'))
word2idx, embedding_matrix = datas[0], datas[1]
else:
word2idx, embedding_matrix = build_dict(train_text, args.embedding_size)
model = Mymodel(args.vocal_size, embedding_matrix, args.hidden_size, args.class_num).to(device)
model.load_state_dict(torch.load('best_model.pth'))
with open(args.data_dif + '\\class.txt', 'r') as f:
class_list = f.read().strip().split('\n')
result = {}
for text in texts:
text_id = [word2idx.get(text_i, 1) for text_i in text]
text_id += [0] * (len_max - len(text_id))
text_id = torch.tensor(text_id).unsqueeze(dim=0).to(device)
model.eval()
with torch.no_grad():
pred = model(text_id)
pred = pred.argmax(dim=1).cpu().numpy().tolist()
true_word = class_list[pred[0]]
result[text] = true_word
print(result)
def main(args):
train_text, train_label, len_max = read_data('train')
if os.path.exists(args.emb_words):
datas = pkl.load(open(args.emb_words, 'rb'))
word2idx, embedding_matrix = datas[0], datas[1]
else:
word2idx, embedding_matrix = build_dict(train_text, args.embedding_size)
train_data = DataLoader(dataset=Datasets(train_text, train_label, word2idx, len_max), batch_size=args.batch_size,
shuffle=True)
dev_text, dev_label, _ = read_data('dev')
dev_data = DataLoader(dataset=Datasets(dev_text, dev_label, word2idx, len_max), batch_size=args.batch_size,
shuffle=True)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = Mymodel(args.vocal_size, embedding_matrix, args.hidden_size, args.class_num).to(device)
loss_func = nn.CrossEntropyLoss()
optimizer = Adam(model.parameters(), lr=args.learning_rate)
best_acc = float("-inf")
for epoch in range(args.epoch):
len_loss, counts = 0, 0
train_pre, train_true = [], []
model.train()
for step, (text, label) in enumerate(train_data):
label = label.to(device)
optimizer.zero_grad()
out = model(text)
loss = loss_func(out, label)
loss.backward()
optimizer.step()
len_loss += loss.item()
counts += 1
train_pre.extend(out.argmax(dim=1).cpu().numpy().tolist())
train_true.extend(label.cpu().numpy().tolist())
if step % 100 == 0:
train_acc = accuracy_score(train_pre, train_true)
print('Epoch:{0:} step:{1} loss:{2:.4f} train_acc:{3:.4f}'.format(epoch, step, len_loss / counts,
train_acc))
model.eval()
dev_pre, dev_true = [], []
for step, (text, label) in enumerate(dev_data):
label = label.to(device)
with torch.no_grad():
out = model(text)
pre = out.argmax(dim=1).cpu().numpy().tolist()
label = label.cpu().numpy().tolist()
dev_pre.extend(pre), dev_true.extend(label)
dev_acc = accuracy_score(dev_pre, dev_true)
print('Epoch:{0} dev_acc:{1:.4f}'.format(epoch, dev_acc))
if best_acc < dev_acc:
best_acc = dev_acc
torch.save(model.state_dict(), 'best_model.pth')
print("训练结束!开始测试")
test_text, test_label, _ = read_data('test')
test_data = DataLoader(dataset=Datasets(test_text, test_label, word2idx, len_max), batch_size=args.batch_size,
shuffle=True)
model = Mymodel(args.vocal_size, embedding_matrix, args.hidden_size, args.class_num).to(device)
model.load_state_dict(torch.load('best_model.pth'))
dev_pre, dev_true = [], []
for step, (text, label) in enumerate(test_data):
label = label.to(device)
with torch.no_grad():
out = model(text)
pre = out.argmax(dim=1).cpu().numpy().tolist()
label = label.cpu().numpy().tolist()
dev_pre.extend(pre), dev_true.extend(label)
acc = accuracy_score(dev_pre, dev_true)
print("测试结果为:{:.4f}".format(acc))
if __name__ == '__main__':
parser = argparse.ArgumentParser(description="MLP文本分类超参数设置")
parser.add_argument("--data_dif", type=str, default=os.path.join(path, 'data', 'class_data'))
parser.add_argument("--out_dir", type=str, default=os.path.join(path, "results"))
parser.add_argument("--bset_model", type=str, default=os.path.join(path, "best-model-mlp.pth"))
parser.add_argument("--emb_words", type=str, default=os.path.join(path, "emb_words-mlp.pkl"))
parser.add_argument("--batch_size", type=int, default=32)
parser.add_argument("--epoch", type=int, default=10)
parser.add_argument("--learning_rate", type=float, default=0.001)
parser.add_argument("--vocal_size", type=int, default=10000)
parser.add_argument("--embedding_size", type=int, default=1024)
parser.add_argument("--hidden_size", type=str, default=1024)
parser.add_argument("--class_num", type=str, default=10)
args = parser.parse_args()
main(args)
texts = ["吃饭了吗"]
generate(args,texts)
目前仅专注于NLP的技术学习和分享
感谢大家的关注与支持!