src/scripts/functionsrc.py

# -*- coding: utf-8 -*-
"""
Created on Mon May  1 20:54:14 2023

@author: sanmo
"""
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader, random_split
import pandas as pd
import json
import os
import gc
from distutils.dir_util import copy_tree
import shutil
import argparse
import flair
from flair.embeddings import TokenEmbeddings, WordEmbeddings, StackedEmbeddings, TransformerWordEmbeddings
from torch import nn, tanh, sigmoid, relu, FloatTensor, rand, stack, optim, cuda, softmax, save, device, tensor, int64, no_grad, concat
from flair.data import Sentence


default_path = 0
entities = 0 
tagger_document = 0
embeddings = 0
json_data = 0
train_loader = 0
val_loader = 0
test_loader = 0
cnn = 0
optimizer = 0
criterion = 0
device = 0
test_sentences = 0
best_valid_loss = np.inf

def check_create(path):
    import os
    
    if not (os.path.isdir(path)):
        os.makedirs(path)
        
def str2bool(v):
    if isinstance(v, bool):
       return v
    if v.lower() in ('yes', 'True','true', 't', 'y', '1'):
        return True
    elif v.lower() in ('no', 'False', 'false', 'f', 'n', '0'):
        return False
    else:
        raise argparse.ArgumentTypeError('Boolean value expected.')
        


class MyDataset(Dataset):
    def __init__(self, len_c1=7, len_c2=5, len_c3=11):
        global json_data
        
        def create_vector(c1,sentence):
            #print("Hola mundo")
            if len(c1): c1 = torch.cat([c1,sentence], dim=0)
            else: c1 = sentence
            return c1
        
        def fix_tensor(tensor, size):
            
            
            while tensor.shape[2] < size:
                tensor = torch.cat([tensor,torch.zeros(1,1,1,1024)], dim=2)
                
            tensor = tensor[:,:,:size,:]
            return tensor
        
        data = []
        self.targets = []
        self.c1=[]
        self.h1=[]
        self.c2=[]
        self.h2=[]
        self.c3=[]
        
        tensor_temp = torch.Tensor(json_data['flat_emb'])
        data = tensor_temp.reshape((tensor_temp.shape[0],1,-1,1024))
        
        self.targets =  create_vector(self.targets,torch.Tensor(json_data['relation']))
        for n_sen in range(tensor_temp.shape[0]):


            tensor_temp = data[n_sen,0,:json_data['h_pos'][n_sen][0],:].reshape((1, 1,-1,1024))
            self.c1 = create_vector(self.c1,fix_tensor(tensor_temp, len_c1))
            
            tensor_temp = data[n_sen,0,json_data['h_pos'][n_sen][0]:json_data['h_pos'][n_sen][-1]+1,:].mean(dim=0).reshape((1,1024))
            self.h1 = create_vector(self.h1,tensor_temp)
            
            tensor_temp = data[n_sen,0,json_data['h_pos'][n_sen][-1]+1:json_data['t_pos'][n_sen][0],:].reshape((1,1,-1,1024))
            self.c2 = create_vector(self.c2,fix_tensor(tensor_temp, len_c2))
            
            tensor_temp = data[n_sen,0,json_data['t_pos'][n_sen][0]:json_data['t_pos'][n_sen][-1]+1,:].mean(dim=0).reshape((1,1024))
            self.h2 = create_vector(self.h2,tensor_temp)                  
            
            tensor_temp = data[n_sen,0,json_data['t_pos'][n_sen][-1]+1:,:].reshape((1, 1,-1,1024))
            self.c3 = create_vector(self.c3,fix_tensor(tensor_temp, len_c3))
        del data
        del tensor_temp
        del json_data
        gc.collect()
        self.targets = self.targets.to(torch.int64)
        #print('Dataset class')
        
    def __len__(self):
        return len(self.targets)

    def __getitem__(self, index):
        c1x = self.c1[index]
        h1x = self.h1[index]
        c2x = self.c2[index]
        h2x = self.h2[index]
        c3x = self.c3[index]
        y = self.targets[index]
        return c1x,h1x,c2x,h2x,c3x, y
    
    
def update_step(c1, h1,c2,h2,c3, label):
    global cnn
    global optimizer
    global criterion
    prediction = cnn(c1, h1,c2,h2,c3)
    optimizer.zero_grad()
    loss = criterion(prediction, label)
    loss.backward()
    optimizer.step()
    acc = (nn.Softmax(dim=1)(prediction).detach().argmax(dim=1) == label).type(torch.float).sum().item()
    #print(acc)
    return loss.item(), acc

def evaluate_step(c1, h1,c2,h2,c3, label):
    global cnn
    global optimizer
    global criterion
    prediction = cnn(c1, h1,c2,h2,c3)
    loss = criterion(prediction, label)
    acc = (nn.Softmax(dim=1)(prediction).detach().argmax(dim=1) == label).type(torch.float).sum().item()
    return loss.item(), acc

def train_one_epoch(epoch, name, rel2id_file):    
    global train_loader
    global val_loader
    global device
    global best_valid_loss
    global optimizer
    global cnn
    global default_path 
    if (device == torch.device('cuda:0')): cnn.cuda()
    
    train_loss, valid_loss, acc_train, acc_valid = 0.0, 0.0, 0.0, 0.0    
    for batch_idx, (c1, h1,c2,h2,c3, targets) in enumerate(train_loader):
        train_loss_temp, acc_train_temp = update_step(c1.to(device), h1.to(device),c2.to(device),h2.to(device),c3.to(device), targets.to(device))  
        train_loss += train_loss_temp
        acc_train += acc_train_temp
    for batch_idx, (c1, h1,c2,h2,c3, targets) in enumerate(val_loader):
        valid_loss_temp, acc_valid_temp = evaluate_step(c1.to(device), h1.to(device),c2.to(device),h2.to(device),c3.to(device), targets.to(device))
        valid_loss += valid_loss_temp
        acc_valid += acc_valid_temp
    # Guardar modelo si es el mejor hasta ahora

    if epoch % 10 == 0:
        # path_save = os.path.normpath(default_path +'/../../models/RC/{}/best_model.pt'.format(name))
        # path_save = path_save.replace('\\', '/')
        # print(os.path.abspath(__file__))
        if valid_loss < best_valid_loss:
            best_valid_loss = valid_loss
            torch.save({'epoca': epoch,
                        'model_state_dict': cnn.state_dict(),
                        'optimizer_state_dict': optimizer.state_dict(),
                        'loss': valid_loss}, 
                         '../../models/RC/{}/best_model.pt'.format(name))
            
            #path_files = default_path + '/../../data/RC/'
            a=0
            #rel2id_file = path_files + 'rel2id.json' 
            
    return train_loss/len(train_loader.dataset), valid_loss/len(val_loader.dataset), acc_train/len(train_loader.dataset), acc_valid/len(val_loader.dataset)


def FocalLoss(input, target, gamma=0, alpha=None, size_average=True):
    from torch.autograd import Variable
    if input.dim()>2:
        input = input.view(input.size(0),input.size(1),-1)  # N,C,H,W => N,C,H*W
        input = input.transpose(1,2)    # N,C,H*W => N,H*W,C
        input = input.contiguous().view(-1,input.size(2))   # N,H*W,C => N*H*W,C
    target = target.view(-1,1)

    logpt = nn.functional.log_softmax(input)
    logpt = logpt.gather(1,target)
    logpt = logpt.view(-1)
    pt = Variable(logpt.data.exp())

    if alpha is not None:
        if alpha.type()!=input.data.type():
            alpha = alpha.type_as(input.data)
        at = alpha.gather(0,target.data.view(-1))
        logpt = logpt * Variable(at)

    loss = -1 * (1-pt)**gamma * logpt
    if size_average: return loss.mean()
    else: return loss.sum()

    
class EarlyStopping:
    def __init__(self, patience=5, min_delta=0):

        self.patience = patience
        self.min_delta = min_delta
        self.counter = 0
        self.min_validation_loss = np.inf
        self.early_stop = False

    def __call__(self, validation_loss):
        if validation_loss < self.min_validation_loss:
            self.min_validation_loss = validation_loss
            self.counter = 0
            self.early_stop = False
            
        elif validation_loss > (self.min_validation_loss + self.min_delta):
            print('Less')
            self.counter += 1
            if self.counter >= self.patience:
                self.early_stop = True
        

                


def SoftmaxModified(x):
    input_softmax = x.transpose(0,1)
    function_activation = nn.Softmax(dim=1)
    output = function_activation(input_softmax)
    output = output.transpose(0,1)
    return output


class MultiModalGMUAdapted(nn.Module):

    def __init__(self, input_size_array, hidden_size, dropoutProbability):
        """Initialize params."""
        super(MultiModalGMUAdapted, self).__init__()
        self.input_size_array = input_size_array
        self.hidden_size = hidden_size
        self.dropout = nn.Dropout(dropoutProbability)
        
        self.h_1_layer = nn.Linear(input_size_array[0], hidden_size, bias=False)
        self.h_2_layer = nn.Linear(input_size_array[1], hidden_size, bias=False)
        self.h_3_layer = nn.Linear(input_size_array[2], hidden_size, bias=False)
        self.h_4_layer = nn.Linear(input_size_array[3], hidden_size, bias=False)
        self.h_5_layer = nn.Linear(input_size_array[4], hidden_size, bias=False)
        
        self.z_1_layer = nn.Linear(input_size_array[0], hidden_size, bias=False)
        self.z_2_layer = nn.Linear(input_size_array[1], hidden_size, bias=False)
        self.z_3_layer = nn.Linear(input_size_array[2], hidden_size, bias=False)
        self.z_4_layer = nn.Linear(input_size_array[3], hidden_size, bias=False)
        self.z_5_layer = nn.Linear(input_size_array[4], hidden_size, bias=False)
        
        
        #self.z_weights = [nn.Linear(input_size_array[m], hidden_size, bias=False) for m in range(modalities_number)]
        #self.input_weights = [nn.Linear(size, hidden_size, bias=False) for size in input_size_array]

        
    def forward(self, inputModalities):
        """Propogate input through the network."""
        # h_modalities = [self.dropout(self.input_weights[i](i_mod)) for i,i_mod in enumerate(inputModalities)]
        # h_modalities = [tanh(h) for h in h_modalities]
        
        h1 = tanh(self.dropout(self.h_1_layer(inputModalities[0])))
        h2 = tanh(self.dropout(self.h_2_layer(inputModalities[1])))
        h3 = tanh(self.dropout(self.h_3_layer(inputModalities[2])))
        h4 = tanh(self.dropout(self.h_4_layer(inputModalities[3])))
        h5 = tanh(self.dropout(self.h_5_layer(inputModalities[4])))

        z1 = self.dropout(self.z_1_layer(inputModalities[0]))
        z2 = self.dropout(self.z_2_layer(inputModalities[1]))
        z3 = self.dropout(self.z_3_layer(inputModalities[2]))   
        z4 = self.dropout(self.z_4_layer(inputModalities[3]))
        z5 = self.dropout(self.z_5_layer(inputModalities[4]))
        

        #z_modalities = [self.dropout(self.z_weights[i](i_mod)) for i,i_mod in enumerate(inputModalities)]
        z_modalities = stack([z1, z2, z3, z4, z5])
        z_normalized = SoftmaxModified(z_modalities)
        final = z_normalized[0] * h1 + z_normalized[1] * h2 + z_normalized[2] * h3 + z_normalized[3] * h4 + z_normalized[4] * h5
        
    
        return final
    
class MyCNN(nn.Module):
    def __init__(self, num_classes=10, len_c1=7, len_c2=5, len_c3=11):
        super(MyCNN, self).__init__()
        shape1 = (((len_c1-2)))#-2)#//2)-2)//2)
        shape2 = (((len_c2-2)))#-2)#//2)-2)//2)
        shape3 = (((len_c3-2)))#-2)#//2)-2)//2)

        # Define convolutional layers
        self.conv_layers1 = nn.Sequential(
            nn.Conv2d(in_channels=1, out_channels=1, kernel_size=(3,1)),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=(shape1,1)),
        )
        
        self.conv_layers2 = nn.Sequential(
            nn.Conv2d(in_channels=1, out_channels=1, kernel_size=(3,1)),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=(shape2,1)),
        )
        
        self.conv_layers3 = nn.Sequential(
            nn.Conv2d(in_channels=1, out_channels=1, kernel_size=(3,1)),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=(shape3,1)),
        )
        

        self.multi_gmu = MultiModalGMUAdapted([1024,1024,1024,1024,1024], 1024, 0.5)

        


        
        self.fc_simple_layers_multi = nn.Sequential(
            nn.Linear(1024 , 256),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(256, num_classes)
        )


    def forward(self, c1, h1,c2,h2,c3):

        # Pass inputs through convolutional layers

        c1 = self.conv_layers1(c1)
        c2 = self.conv_layers2(c2)
        c3 = self.conv_layers3(c3)
        #print(c1.shape)
        
        h1 = tanh(h1)
        h2 = tanh(h2)
        #print(c1.shape)
        c1 = torch.flatten(c1, start_dim=1)
        c2 = torch.flatten(c2, start_dim=1)
        c3 = torch.flatten(c3, start_dim=1)
        #print(c1.shape)
        
        
        # Multi GMU
        mgmu_out = self.multi_gmu([c1,h1,c2,h2,c3]) 
        x = self.fc_simple_layers_multi(mgmu_out)
        

        # Return final output
        return x
    
def define_model():
    global cnn
    global optimizer
    global criterion
    
    cnn = MyCNN()
    optimizer = torch.optim.Adam(cnn.parameters(), lr=0.001)
    criterion = lambda pred,tar: FocalLoss(input=pred,target=tar,gamma=0.7)
    
def train_model(name, epocs, rel2id_path):
    max_epochs, best_valid_loss = epocs, np.inf
    running_loss = np.zeros(shape=(max_epochs, 4))
    early_stopping = EarlyStopping(patience=10, min_delta=0.01)
    
    for epoch in range(max_epochs):
        running_loss[epoch] = train_one_epoch(epoch, name, rel2id_path)
        early_stopping(running_loss[epoch, 1])
        print(f"Epoch {epoch} \t Train_loss = {running_loss[epoch, 0]:.4f} \t Valid_loss = {running_loss[epoch, 1]:.4f} \n\t\t\t Train_acc = {running_loss[epoch, 2]:.4f} \t Valid_acc = {running_loss[epoch, 3]:.4f}")
        if early_stopping.early_stop:
          print("We are at epoch:", epoch)
          break
      
        
def usage_cuda_rc(cuda):
    global device
    if cuda:
        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        flair.device = device
        if flair.device == torch.device('cpu'): return 'Error handling GPU, CPU will be used'
        elif flair.device == torch.device('cuda:0'): return 'GPU detected, GPU will be used'
    else:
        device = torch.device('cpu')
        flair.device = device
        return 'CPU will be used'
    
    
def create_embbedings():
    global embeddings
    if (not embeddings):
        embeddings = TransformerWordEmbeddings(
                    model='xlm-roberta-large',
                    layers="-1",
                    subtoken_pooling="first", 
                    fine_tune=True,
                    use_context=True,
                )



    
def prepare_data(rel2id_path, path_data):
    create_embbedings()
    global embeddings
    global json_data
    #Embbeb data
    
    global default_path 
    
    
    #path_files 
    
    
    #rel2id_file = path_files + 'rel2id.json' 
    #shutil.copy(rel2id_path, rel2id_file)   
    with open(rel2id_path, mode='r') as f: 
        rel2id = json.load(f)
        
    #path_data = path_files+"train.txt"

    #Json to save the data
    json_data = {'flat_emb':[], 'relation':[], 'h_pos':[], 't_pos':[]}
    PADDING = np.zeros(1024)
    doc=0
    with open(path_data, mode='r', encoding='utf-8') as f:
        sentence_temp = []
        h_pos = []
        t_pos = []
        current_ent=''
        cont=0
        
        for n,line in enumerate(f.readlines()):
            if line != '\n':
                sentence_temp.append(line.split('\t')[0])
                
                if line.split('\t')[1] != 'O':
                    if current_ent == '':
                        h_pos.append(cont)
                        current_ent = line.split('\t')[1]
                        
                    elif line.split('\t')[1] == current_ent:
                        h_pos.append(cont)
                        
                    else:
                        t_pos.append(cont)
                        
                if line.split('\t')[2].replace('\n','') != '-' : relation = line.split('\t')[2].replace('\n','') 

                cont += 1
                
            else:
                
                #Embbedding sentence
                sentence = Sentence(sentence_temp)
                embeddings.embed(sentence)
                
                

                sentence_emb_flatten = []
                for tk in sentence: 
                    #flatten_embeddings    
                    if len(sentence_emb_flatten): sentence_emb_flatten = np.hstack((sentence_emb_flatten,
                                                                               tk.embedding.detach().to('cpu').numpy()))
                    else: sentence_emb_flatten = tk.embedding.detach().to('cpu').numpy()
                       
                number_padding = 100 - len(sentence)
                
                if number_padding > 0: 
                    for pd in range(number_padding):
                        sentence_emb_flatten = np.hstack((sentence_emb_flatten,
                                                          PADDING))
                       
                #Save embeddings information
                json_data['flat_emb'].append(list(sentence_emb_flatten))
                json_data['h_pos'].append(h_pos)
                json_data['t_pos'].append(t_pos)
                json_data['relation'].append(rel2id[relation])
                
                sentence_temp = []
                h_pos = []
                t_pos = []
                current_ent=''
                cont=0
    dataset = MyDataset()

    train_set_size = int(len(dataset) * 0.9)
    valid_set_size = len(dataset) - train_set_size
        
    train_dataset, val_dataset = random_split(dataset, [train_set_size, valid_set_size ])
    del dataset
    global train_loader 
    global val_loader
                                                 
    train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=64, shuffle=True)
    
def prepare_data_test(name, path_data):
    create_embbedings()
    global embeddings
    global json_data
    global test_sentences
    global entities
    #Embbeb data
    
    global default_path 
    
    
    test_sentences = []
    entities = []
    #path_files = default_path + '/../../data/RC/'
    # path_model = default_path + '/../../models/RC/{}/'.format(name)

        
        
    #path_data = path_files+"test.txt"

    #Json to save the data
    json_data = {'flat_emb':[], 'relation':[], 'h_pos':[], 't_pos':[]}
    PADDING = np.zeros(1024)
    doc=0
   
    with open(path_data, mode='r', encoding='utf-8') as f:
         
        sentence_temp = []
        entities_temp = []
        h_pos = []
        t_pos = []
        current_ent=''
        cont=0
        for n,line in enumerate(f.readlines()):
            #print(line)
            if line != '\n':
                sentence_temp.append(line.split('\t')[0])
                entities_temp.append(line.split('\t')[1])
                if line.split('\t')[1] != 'O':
                    if current_ent == '':
                        h_pos.append(cont)
                        current_ent = line.split('\t')[1]
                        
                    elif line.split('\t')[1] == current_ent:
                        h_pos.append(cont)
                        
                    else:
                        t_pos.append(cont)
                        
              #  if line.split('\t')[2].replace('\n','') != '-' : relation = '-'

                cont += 1
                
            else:
                
                #Embbedding sentence
                sentence = Sentence(sentence_temp)
                
                test_sentences.append(sentence_temp)
                entities.append(entities_temp)
                #print('mid while')
                embeddings.embed(sentence)
                

                sentence_emb_flatten = []
                for tk in sentence: 
                    #flatten_embeddings    
                    if len(sentence_emb_flatten): sentence_emb_flatten = np.hstack((sentence_emb_flatten,
                                                                               tk.embedding.detach().to('cpu').numpy()))
                    else: sentence_emb_flatten = tk.embedding.detach().to('cpu').numpy()
                       
                number_padding = 100 - len(sentence)
                
                if number_padding > 0: 
                    for pd in range(number_padding):
                        sentence_emb_flatten = np.hstack((sentence_emb_flatten,
                                                          PADDING))
                       
                #Save embeddings information
                json_data['flat_emb'].append(list(sentence_emb_flatten))
                json_data['h_pos'].append(h_pos)
                json_data['t_pos'].append(t_pos)
                json_data['relation'].append(1)
                
                sentence_temp = []
                entities_temp = []
                h_pos = []
                t_pos = []
                current_ent=''
                cont=0

                
    dataset = MyDataset()
    global test_loader 
    test_loader = DataLoader(dataset, batch_size=64, shuffle=False)
        

    del dataset
    

                                                 

#------------------Backend functions----------------------------------------

def training_model_rc(name, path_data, rel2id_path, epochs=200):
    global default_path 
    default_path = os.path.dirname(os.path.abspath(__file__))
    default_path = default_path.replace('\\', '/')
    print(name)
    #FUNCION
    
    try:
        define_model()
    except:
        return 13
    print('Model defined')
    check_create(default_path + '/../../models/RC/{}/'.format(name))
    
    try:
        prepare_data(rel2id_path, path_data)
    except:
        return 12

    print('Data prepared')
    #Train the model
    try:
        train_model(name, epochs, rel2id_path)
    except:
        return 7
    #save the model in
    path_model = default_path + '/../../models/RC/{}/best_model.pt'.format(name)
    
    shutil.copy(rel2id_path, default_path + '/../../models/RC/{}/rel2id.json'.format(name))

    return "model trined and saved at {}".format(path_model)

    
def use_model_rc(name, path_data, output_dir):
    global default_path 
    default_path = os.path.dirname(os.path.abspath(__file__))
    default_path = default_path.replace('\\', '/')
    #--------------Load the trained model-------------------------
    path_model = default_path + '/../../models/RC/{}/best_model.pt'.format(name)

    
    rel2id_file = default_path + '/../../models/RC/{}/rel2id.json'.format(name)
    with open(rel2id_file, mode='r') as f: 
        rel2id = json.load(f)
    id2rel = [m for _,m in sorted(zip(list(rel2id.values()),list(rel2id.keys())), key=lambda pair: pair[0])]
    
    if not (os.path.isfile(path_model)): 
        print('Model does not exists')
        return 10
    print(path_data)
    if not os.path.isfile(path_data): 
        print('Input file is not a file')
        return 9 
    
    global cnn
    try:
        cnn = MyCNN()
    except:
        return 13
    
    print('Model defined')
    
    try:
        saved_model = torch.load(path_model)
        cnn.load_state_dict(saved_model['model_state_dict'])
    except:
        return 1
    print('Model loaded')
    
    #-----------------Load the document-------------------------
    try:
        prepare_data_test(name, path_data)
    except:
        return 12
    
    global json_data
    print('Data prepared')
    #-----------------Predict-------------------------
    
    global test_loader
    ypred = []
    relations = []
    for batch_idx, (c1, h1,c2,h2,c3, targets)  in enumerate(test_loader):
        x = cnn(c1, h1,c2,h2,c3)
        ypred.append(nn.Softmax(dim=1)(x).detach().argmax(dim=1))
    ypred = np.concatenate(ypred)

    relations = [id2rel[rel] for rel in ypred]
    print('prediction')
    #-----------------Tagged the document-------------------------
    global test_sentences
    global entities
    results = {'sentences':{'tokens':test_sentences, 'entities':entities}, 'relations': relations}
    

    #-----------------Save the results-------------------------
    try:
        with open(output_dir, "w", encoding='utf-8') as write_file:
            json.dump(results, write_file)
    
        print('-'*20,'Tagged complete','-'*20)
        print('Document tagged saved in {}'.format(output_dir))
    except:
        print('Error in output file')
        return 11
       
    return results