-
Notifications
You must be signed in to change notification settings - Fork 0
/
rnn_predict.py
223 lines (175 loc) · 8.02 KB
/
rnn_predict.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import math
from torchtext import data
from torchtext import vocab
import jieba
class Config(object):
"""配置参数"""
def __init__(self, embedding, type):
self.clip = 10
self.model_name = 'rnn'
self.class_list = ['教育', '财经', '时政', '科技', '社会', '健康', '其他']
# self.vocab_path = dataset + '/data/vocab.pkl' # 词表
self.save_path = './models/' + self.model_name + '.pt' # 模型训练结果
self.log_path = './logs/' + self.model_name
# embedding size
self.embedding_file = embedding
self.embed = 300
# 预训练词向量
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # 设备
self.dropout = 0.3 # 随机失活
self.require_improvement = 1000 # 若超过 1000 batch效果还没提升，则提前结束训练
self.num_classes = len(self.class_list) # 类别数
self.num_epochs = 60 # epoch数
self.bidirectional = True
self.pad_size = 512 # 每句话处理成的长度(短填长切)
self.learning_rate = 1e-2 # 学习率
self.hidden_size = 128 # lstm隐藏层
self.num_layers = 2 # lstm层数
if type == 'test':
self.batch_size = 16 # mini-batch 大小 128
self.n_vocab = 1000 # 词表大小，在运行时赋值
self.data_path = './data/test/'
self.print_step = 2 # 100
else:
self.batch_size = 128 # mini-batch 大小 128
self.n_vocab = 15000 # 15000 词表大小，在运行时赋值
self.data_path = './fold/' # './fold/'
self.print_step = 100 # 100
class News_clf(nn.Module):
def __init__(self, config, word_emb):
super(News_clf, self).__init__()
self.word_embedding = nn.Embedding.from_pretrained(word_emb)
self.lstm = nn.LSTM(input_size=config.embed, hidden_size=config.hidden_size, num_layers=config.num_layers,
bidirectional=config.bidirectional, batch_first=True, dropout=config.dropout)
self.w = nn.Parameter(torch.zeros(config.hidden_size * 2))
self.dropout = nn.Dropout(config.dropout)
if config.bidirectional:
input_features = config.hidden_size * 2
else:
input_features = config.hidden_size
self.fc = nn.Linear(input_features, config.num_classes)
def attention_net(self, x, query, mask=None): # 软性注意力机制（key=value=x）
dv = query.size(1)
d_k = query.size(-1) # d_k为query的维度
scores = torch.matmul(query, x.transpose(1, 2)) / math.sqrt(d_k) # 打分机制 scores:[batch, seq_len, seq_len]
p_attn = F.softmax(scores, dim=-1) # 对最后一个维度归一化得分
context = torch.matmul(p_attn, x).sum(1) / dv # 对权重化的x求和，[batch, seq_len, hidden_dim*2]->[batch, hidden_dim*2]
return context, p_attn
def forward(self, x):
text, _ = x
# title_emb: [seq_len, batch_size, emd_dim]
text_emb = self.dropout(self.word_embedding(text))
# [batch_size, emd_dim, seq_len]
text_emb = text_emb.permute(1, 0, 2)
H, _ = self.lstm(text_emb) # [batch_size, seq_len, hidden_size * num_direction]
query = self.dropout(H)
attn_output, attention = self.attention_net(H, query)
out = self.fc(attn_output)
return out
# load data
def load_news(config, text_field, band_field):
fields = {
'text': ('text', text_field),
'label': ('label', band_field)
}
word_vectors = vocab.Vectors(config.embedding_file)
train, val, test = data.TabularDataset.splits(
path=config.data_path, train='train.csv', validation='val.csv',
test='test.csv', format='csv', fields=fields)
print("the size of train: {}, dev:{}, test:{}".format(
len(train.examples), len(val.examples), len(test.examples)))
text_field.build_vocab(train, val, test, max_size=config.n_vocab, vectors=word_vectors,
unk_init=torch.Tensor.normal_)
train_iter, val_iter, test_iter = data.BucketIterator.splits(
(train, val, test), batch_sizes=(config.batch_size, config.batch_size, config.batch_size), sort=False,
device=config.device, sort_within_batch=False, shuffle=False)
return train_iter, val_iter, test_iter
embedding = './data/sgns.sogou.word'
type = 'train'
config = Config(embedding, type)
print('device:', config.device)
def tokenizer(s):
return jieba.lcut(s)
# data loader and split Chinese
text_field = data.Field(tokenize=tokenizer, include_lengths=True, fix_length=config.pad_size)
band_field = data.Field(sequential=False, use_vocab=False, batch_first=True,
dtype=torch.int64, preprocessing=data.Pipeline(lambda x: int(x)))
train_iterator, val_iterator, test_iterator = load_news(config, text_field, band_field)
# initialize
def count_parameters(model):
return sum(p.numel() for p in model.parameters() if p.requires_grad)
word_emb = text_field.vocab.vectors
model = News_clf(config, word_emb)
model = model.to(config.device)
dict = torch.load(config.save_path, map_location=config.device)
model.load_state_dict(dict['model_state_dict'])
model.eval()
# from sklearn import metrics
# def classifiction_metric(preds, labels, label_list):
# acc = metrics.accuracy_score(labels, preds)
# labels_list = [i for i in range(len(label_list))]
# report = metrics.classification_report(
# labels, preds, labels=labels_list, target_names=label_list, digits=4, output_dict=True)
# return acc, report
# def evaluation(model, iterator, config):
# model.eval()
# all_preds = np.array([], dtype=int)
# all_labels = np.array([], dtype=int)
# with torch.no_grad():
# for batch in iterator:
# logits = model(batch.text)
# label = batch.label.detach().cpu().numpy()
#
# logits = F.softmax(logits, dim=1)
#
# preds = logits.detach().cpu().numpy()
#
# preds = np.argmax(preds, axis=1)
#
# all_preds = np.append(all_preds, preds)
# all_labels = np.append(all_labels, label)
# return classifiction_metric(all_preds, all_labels, config.class_list)
# evaluation
# test_acc, test_report = evaluation(model, test_iterator, config)
# print(test_acc, test_report)
class News():
def __init__(self, s):
self.text = (s.unsqueeze(1), torch.tensor([len(s)]).to(config.device))
def predict_label(t, c):
input_doc = tokenizer(t + c)
input_doc = input_doc + ['<pad>'] * (config.pad_size - len(input_doc)) if len(
input_doc) < config.pad_size else input_doc[:config.pad_size]
indexed = [text_field.vocab.stoi[t] for t in input_doc]
with torch.no_grad():
res = model(News(torch.tensor(indexed).to(config.device)).text)
# print(res.detach().cpu().numpy()[0])
res = F.softmax(res, dim=1)
resp = res.detach().cpu().numpy()[0]
index = np.argmax(resp)
if resp[index] > 0.5:
print('label:', config.class_list[index])
return config.class_list[index]
return config.class_list[-1]
# import pandas as pd
# df = pd.read_csv('./data/sample.csv')
# df['predict'] = None
# for i in range(len(df)):
# tmp = str(df['title'][i]) # if type(df['title'][i])
# c = str(df['content'][i]) # if type(df['content'][i]) == str else ''
# input_doc = tokenizer(tmp + c)
# input_doc = input_doc + ['<pad>'] * (config.pad_size - len(input_doc)) if len(
# input_doc) < config.pad_size else input_doc[:config.pad_size]
# indexed = [text_field.vocab.stoi[t] for t in input_doc]
#
# with torch.no_grad():
# res = model(News(torch.tensor(indexed).to(config.device)).text)
# resp = res.detach().cpu().numpy()[0]
# print('value: ', resp)
# index = np.argmax(resp)
# df['predict'][i] = config.class_list[index]
#
# df.to_csv('predict.csv', index=False)