forked from kubapok/en-ner-conll-2003
124 lines
4.1 KiB
Python
124 lines
4.1 KiB
Python
import torch
|
|
import pandas as pd
|
|
from torchtext.vocab import Vocab
|
|
from collections import Counter
|
|
|
|
x_train = pd.read_table('train/train.tsv', sep='\t', header = None)
|
|
x_dev = pd.read_table('dev-0/in.tsv', sep='\t', header = None)
|
|
y_dev = pd.read_table('dev-0/expected.tsv', sep='\t', header = None)
|
|
x_test = pd.read_table('test-A/in.tsv', sep='\t', header = None)
|
|
|
|
label_list = ['O', 'B-LOC', 'I-LOC', 'B-MISC', 'I-MISC', 'B-ORG', 'I-ORG', 'B-PER', 'I-PER']
|
|
x_train[0] = x_train[0].apply(lambda x: [label_list.index(i) for i in x.split()])
|
|
x_train[1] = x_train[1].apply(lambda x: x.split())
|
|
|
|
x_dev[0] = x_dev[0].apply(lambda x: x.split())
|
|
x_test[0] = x_test[0].apply(lambda x: x.split())
|
|
|
|
def data_process(dt):
|
|
return [ torch.tensor([vocab['<bos>']] +[vocab[token] for token in document ] + [vocab['<eos>']], dtype = torch.long) for document in dt]
|
|
|
|
def build_vocab(dataset):
|
|
counter = Counter()
|
|
for document in dataset:
|
|
counter.update(document)
|
|
return Vocab(counter, specials=['<unk>', '<pad>', '<bos>', '<eos>'])
|
|
|
|
vocab = build_vocab(x_train[1])
|
|
|
|
def labels_process(dt):
|
|
return [ torch.tensor([0] + document + [0], dtype = torch.long) for document in dt]
|
|
|
|
class NERModel(torch.nn.Module):
|
|
|
|
def __init__(self,):
|
|
super(NERModel, self).__init__()
|
|
self.emb = torch.nn.Embedding(23627,200)
|
|
self.fc1 = torch.nn.Linear(600,9)
|
|
|
|
|
|
def forward(self, x):
|
|
x = self.emb(x)
|
|
x = x.reshape(600)
|
|
x = self.fc1(x)
|
|
return x
|
|
|
|
ner_model = NERModel()
|
|
criterion = torch.nn.CrossEntropyLoss()
|
|
optimizer = torch.optim.Adam(ner_model.parameters())
|
|
|
|
train_labels = labels_process(x_train[0])
|
|
train_tokens_ids = data_process(x_train[1])
|
|
|
|
for epoch in range(2):
|
|
loss_score = 0
|
|
acc_score = 0
|
|
prec_score = 0
|
|
selected_items = 0
|
|
recall_score = 0
|
|
relevant_items = 0
|
|
items_total = 0
|
|
ner_model.train()
|
|
for i in range(len(train_labels)):
|
|
for j in range(1, len(train_labels[i])-1):
|
|
X = train_tokens_ids[i][j-1: j+2]
|
|
Y = train_labels[i][j: j+1]
|
|
|
|
Y_predictions = ner_model(X)
|
|
|
|
acc_score += int(torch.argmax(Y_predictions) == Y)
|
|
|
|
if torch.argmax(Y_predictions) != 0:
|
|
selected_items +=1
|
|
if torch.argmax(Y_predictions) != 0 and torch.argmax(Y_predictions) == Y.item():
|
|
prec_score +=1
|
|
|
|
if Y.item() != 0:
|
|
relevant_items += 1
|
|
if Y.item() != 0 and torch.argmax(Y_predictions) == Y.item():
|
|
recall_score += 1
|
|
|
|
items_total +=1
|
|
optimizer.zero_grad()
|
|
loss = criterion(Y_predictions.unsqueeze(0), Y)
|
|
loss.backward()
|
|
optimizer.step()
|
|
loss_score += loss.item()
|
|
|
|
precision = prec_score / selected_items
|
|
recall = recall_score / relevant_items
|
|
f1_score = (2*precision * recall) / (precision + recall)
|
|
print('epoch: ', epoch)
|
|
print('loss: ', loss_score / items_total)
|
|
print('acc: ', acc_score / items_total)
|
|
print('prec: ', precision)
|
|
print('recall: ', recall)
|
|
print('f1: ', f1_score)
|
|
|
|
dev_data_tokens_ids = data_process(x_dev[0])
|
|
dev_results = []
|
|
for i in range(len(dev_data_tokens_ids)):
|
|
line = []
|
|
for j in range(1, len(dev_data_tokens_ids[i]) - 1):
|
|
X = dev_data_tokens_ids[i][j-1: j+2]
|
|
Y_predictions = ner_model(X)
|
|
result = torch.argmax(Y_predictions)
|
|
label = label_list[result]
|
|
line.append(label)
|
|
dev_results.append(line)
|
|
|
|
test_data_tokens_ids = data_process(x_test[0])
|
|
test_results = []
|
|
for i in range(len(test_data_tokens_ids)):
|
|
line = []
|
|
for j in range(1, len(test_data_tokens_ids[i]) - 1):
|
|
X = test_data_tokens_ids[i][j-1: j+2]
|
|
Y_predictions = ner_model(X)
|
|
result = torch.argmax(Y_predictions)
|
|
label = label_list[result]
|
|
line.append(label)
|
|
test_results.append(line)
|
|
|
|
dev_results.to_csv(r'dev-0/out.tsv', sep='\t', index=False, header=False)
|
|
test_results.to_csv(r'test-A/out.tsv', sep='\t', index=False, header=False)
|