en-ner-conll-2003/ner-rnn.py

import torch
from torchtext.vocab import Vocab
import lzma
from torchcrf import CRF
import numpy as np
import re
import itertools
from datasets import load_dataset
from torchtext.vocab import Vocab
from collections import Counter

# zadanie domowe oparte na jupyter notebook z zajec https://git.wmi.amu.edu.pl/filipg/aitech-eks-pub/src/branch/master/cw/11_NER_RNN_ODPOWIEDZI.ipynb
# i https://git.wmi.amu.edu.pl/filipg/aitech-eks-pub/src/branch/master/cw/09_sequence_labeling_ODPOWIEDZI.ipynb

# model sieci GRU
class GRU(torch.nn.Module):

    def __init__(self):
        super(GRU, self).__init__()
        self.emb = torch.nn.Embedding(len(vocab.itos), 100)
        self.dropout = torch.nn.Dropout(0.2)
        self.rec = torch.nn.GRU(100, 256, 2, batch_first = True, bidirectional = True)
        self.fc1 = torch.nn.Linear(2* 256 , 9)
        
    def forward(self, x):
        emb = torch.relu(self.emb(x))
        emb = self.dropout(emb)
        gru_output, h_n = self.rec(emb)
        out_weights = self.fc1(gru_output)
        return out_weights

	    
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

def read_data(filename):
    data = lzma.open(filename).read().decode('UTF-8').split('\n')
    return [line.split('\t') for line in data][:-1]

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 get_scores(y_true, y_pred):
     acc_score = 0
     tp = 0
     selected_items = 0
     relevant_items = 0 
 
     for p,t in zip(y_pred, y_true):
         if p == t:
             acc_score +=1
 
         if p > 0 and p == t:
             tp +=1
 
         if p > 0:
             selected_items += 1
 
         if t > 0 :
             relevant_items +=1
             
     if selected_items == 0:
         precision = 1.0
     else:
         precision = tp / selected_items
             
     if relevant_items == 0:
         recall = 1.0
     else:
         recall = tp / relevant_items
     
     if precision + recall == 0.0 :
         f1 = 0.0
     else:
         f1 = 2* precision * recall  / (precision + recall)
 
     return precision, recall, f1

def labels_process(dt):
    return [ torch.tensor([0] + document + [0], 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>'])

def predict_labels(dataset_tokens, dataset_labels, model):
    print(len(dataset_tokens[0]), len(dataset_labels[0]))
    Y_true = []
    Y_pred = []
    result = []
    for i in tqdm(range(len(dataset_labels))):
        batch_tokens = dataset_tokens[i].unsqueeze(1)
        tags = list(dataset_labels[i].numpy())
        emissions = gru(batch_tokens).squeeze(0)
        tmp = crf.decode(emissions)[0]
        Y_pred += tmp
        result += [tmp]
        Y_true += tags
    print(get_scores(Y_true, Y_pred))
    return result

def predict(path, model):
    with open(path + '/in.tsv', "r", encoding="utf-8") as f:
        data = [line.rstrip() for line in f]
    data = [i.split() for i in data]
    tokens_ids = data_process(data)
    
    Y_true = []
    Y_pred = []
    result = []
    for i in tqdm(range(len(tokens_ids))):
        batch_tokens = tokens_ids[i].unsqueeze(1)
        emissions = gru(batch_tokens).squeeze(0)
        tmp = crf.decode(emissions)[0]
        Y_pred += tmp
        result += [tmp]
    r = [[ner_tags_set[i] for i in tmp] for tmp in result]
    r = [i[1:-1] for i in r]
    for doc in r:
        if doc[0] != 'O':
            doc[0] = 'B' + doc[0][1:]
        for i in range(len(doc))[:-1]:
            if doc[i] == 'O':
                if doc[i + 1] != 'O':
                    doc[i + 1] = 'B' + doc[i + 1][1:]
            elif doc[i + 1] != 'O':
                if doc[i][1:] == doc[i + 1][1:]:
                    doc[i + 1] = 'I' + doc[i + 1][1:]
                else:
                    doc[i + 1] = 'B' + doc[i + 1][1:]
    f = open(path + "/out.tsv", "a")
    for i in r:
        f.write(' '.join(i) + '\n')
    f.close()
    return result
     
# uzycie crf
def eval_model(dataset_tokens, dataset_labels, model):
     Y_true = []
     Y_pred = []
     for i in tqdm(range(len(dataset_labels))):
         batch_tokens = dataset_tokens[i].unsqueeze(1)
         tags = list(dataset_labels[i].numpy())
         emissions = gru(batch_tokens).squeeze(0)
         Y_pred += crf.decode(emissions)[0]
         Y_true += tags
     return get_scores(Y_true, Y_pred)


device_gpu = torch.device("cuda:0")
gru = GRU()
criterion = torch.nn.CrossEntropyLoss()
params = list(gru.parameters()) + list(crf.parameters())
optimizer = torch.optim.Adam(gru.parameters())
EPOCHS = 5

train_set = read_data('train/train.tsv.xz')

tokens, ner_tags = [], []
for i in train_set:
    ner_tags.append(i[0].split())
    tokens.append(i[1].split())

vocab = build_vocab(tokens)

train_tokens_ids = data_process(tokens)

ner_tags_set = list(set(itertools.chain(*ner_tags)))
ner_tags_set.sort()
train_labels = labels_process([[ner_tags_set.index(token) for token in doc] for doc in ner_tags])
num_tags = max([max(x) for x in train_labels]) + 1 

for i in range(EPOCHS):
    gru.train()
    crf.train()
    for i in tqdm(range(len(train_labels))):
        batch_tokens = train_tokens_ids[i].unsqueeze(1)
        tags = train_labels[i].unsqueeze(1)
        emissions = gru(batch_tokens).squeeze(0)
        optimizer.zero_grad()
        loss = -crf(emissions,tags.squeeze(0))
        loss.backward()
        optimizer.step()
    gru.eval()
    crf.eval()

with open('dev-0/in.tsv', "r", encoding="utf-8") as f:
    dev_0_data = [line.rstrip() for line in f]
    
dev_0_data = [i.split() for i in dev_0_data]
dev_0_tokens_ids = data_process(dev_0_data)

with open('dev-0/expected.tsv', "r", encoding="utf-8") as f:
     dev_0_labels = [line.rstrip() for line in f]
     
dev_0_labels = [i.split() for i in dev_0_labels]
dev_0_labels = labels_process([[ner_tags_set.index(token) for token in doc] for doc in dev_0_labels])

tmp = predict_labels(dev_0_tokens_ids, dev_0_labels, gru)

r = [[ner_tags_set[i] for i in tmp2] for tmp2 in tmp]
r = [i[1:-1] for i in r]


for doc in r:
    if doc[0] != 'O':
        doc[0] = 'B' + doc[0][1:]
    for i in range(len(doc))[:-1]:
        if doc[i] == 'O':
            if doc[i + 1] != 'O':
                doc[i + 1] = 'B' + doc[i + 1][1:]
        elif doc[i + 1] != 'O':
            if doc[i][1:] == doc[i + 1][1:]:
                doc[i + 1] = 'I' + doc[i + 1][1:]
            else:
                doc[i + 1] = 'B' + doc[i + 1][1:]


predict('dev-0', gru)
predict('test-A', gru)

f = open("dev-0/out.tsv", "a")
for i in r:
    f.write(' '.join(i) + '\n')
f.close()