forked from kubapok/en-ner-conll-2003
104 lines
3.3 KiB
Python
104 lines
3.3 KiB
Python
|
from os import sep
|
||
|
|
from nltk import word_tokenize
|
||
|
|
import pandas as pd
|
||
|
|
import torch
|
||
|
|
from tqdm import tqdm
|
||
|
|
from torchtext.vocab import vocab
|
||
|
|
from collections import Counter, OrderedDict
|
||
|
|
import spacy
|
||
|
|
from torchcrf import CRF
|
||
|
|
from torch.utils.data import DataLoader
|
||
|
|
|
||
|
|
nlp = spacy.load('en_core_web_sm')
|
||
|
|
|
||
|
|
class Model(torch.nn.Module):
|
||
|
|
def __init__(self, num_tags, seq_length):
|
||
|
|
super(Model, self).__init__()
|
||
|
|
self.emb = torch.nn.Embedding(len(vocab.get_itos()), 100)
|
||
|
|
self.gru = torch.nn.GRU(100, 256, 1, batch_first=True)
|
||
|
|
self.hidden2tag = torch.nn.Linear(256, 9)
|
||
|
|
self.crf = CRF(num_tags, batch_first=True)
|
||
|
|
self.relu = torch.nn.ReLU()
|
||
|
|
self.fc1 = torch.nn.Linear(1, seq_length)
|
||
|
|
self.softmax = torch.nn.Softmax(dim=0)
|
||
|
|
self.sigm = torch.nn.Sigmoid()
|
||
|
|
|
||
|
|
def forward(self, data, tags):
|
||
|
|
emb = self.relu(self.emb(data))
|
||
|
|
out, h_n = self.gru(emb)
|
||
|
|
# out = self.dense1(out.squeeze(0).T)
|
||
|
|
out = self.hidden2tag(out)
|
||
|
|
out = self.crf(out, tags.T)
|
||
|
|
out = self.sigm(self.fc1(torch.tensor([out])))
|
||
|
|
return out
|
||
|
|
|
||
|
|
|
||
|
|
def process_document(document):
|
||
|
|
# return [str(tok.lemma) for tok in nlp(document)]
|
||
|
|
return document.split(" ")
|
||
|
|
|
||
|
|
def build_vocab(dataset):
|
||
|
|
counter = Counter()
|
||
|
|
for document in dataset:
|
||
|
|
counter.update(process_document(document))
|
||
|
|
sorted_by_freq_tuples = sorted(counter.items(), key=lambda x: x[1], reverse=True)
|
||
|
|
ordered_dict = OrderedDict(sorted_by_freq_tuples)
|
||
|
|
v = vocab(counter)
|
||
|
|
default_index = -1
|
||
|
|
v.set_default_index(default_index)
|
||
|
|
return v
|
||
|
|
|
||
|
|
def data_process(dt):
|
||
|
|
return [ torch.tensor([vocab[token] for token in document.split(" ") ], dtype = torch.long) for document in dt]
|
||
|
|
|
||
|
|
def labels_process(dt):
|
||
|
|
return [ torch.tensor([labels_vocab[token] for token in document.split(" ") ], dtype = torch.long) for document in dt]
|
||
|
|
|
||
|
|
data = pd.read_csv("train/train.tsv", sep="\t")
|
||
|
|
data.columns = ["labels", "text"]
|
||
|
|
|
||
|
|
vocab = build_vocab(data['text'])
|
||
|
|
# labels_vocab = build_vocab(data['labels'])
|
||
|
|
|
||
|
|
labels_vocab = {
|
||
|
|
'O': 0,
|
||
|
|
'B-PER': 1,
|
||
|
|
'B-LOC': 2,
|
||
|
|
'I-PER': 3,
|
||
|
|
'B-MISC': 4,
|
||
|
|
'I-MISC': 5,
|
||
|
|
'I-LOC': 6,
|
||
|
|
'B-ORG': 7,
|
||
|
|
'I-ORG': 8
|
||
|
|
}
|
||
|
|
|
||
|
|
train_tokens_ids = data_process(data["text"])
|
||
|
|
train_labels = labels_process(data["labels"])
|
||
|
|
num_tags = 9
|
||
|
|
NUM_EPOCHS = 5
|
||
|
|
seq_length = 15
|
||
|
|
|
||
|
|
model = Model(num_tags, seq_length)
|
||
|
|
criterion = torch.nn.CrossEntropyLoss()
|
||
|
|
optimizer = torch.optim.Adam(model.parameters())
|
||
|
|
|
||
|
|
train_dataloader = DataLoader(list(zip(train_tokens_ids, train_labels)), batch_size=64, shuffle=True)
|
||
|
|
# test_dataloader = DataLoader(train_labels, batch_size=64, shuffle=True)
|
||
|
|
|
||
|
|
for i in range(NUM_EPOCHS):
|
||
|
|
model.train()
|
||
|
|
#for i in tqdm(range(500)):
|
||
|
|
for i in tqdm(range(len(train_labels))):
|
||
|
|
for k in range(0, len(train_tokens_ids[i]) - seq_length, seq_length):
|
||
|
|
batch_tokens = train_tokens_ids[i][k: k + seq_length].unsqueeze(0)
|
||
|
|
tags = train_labels[i][k: k + seq_length].unsqueeze(1)
|
||
|
|
|
||
|
|
predicted_tags = model(batch_tokens, tags)
|
||
|
|
|
||
|
|
optimizer.zero_grad()
|
||
|
|
tags = torch.tensor([x[0] for x in tags])
|
||
|
|
loss = criterion(predicted_tags.unsqueeze(0),tags.T)
|
||
|
|
|
||
|
|
loss.backward()
|
||
|
|
optimizer.step()
|
||
|
|
model.zero_grad()
|