Model/MultimodelNER/VLSP2016/train_umt_2016.py

import os
import sys

os.environ["CUDA_VISIBLE_DEVICES"] = "0"
import argparse

import logging
import random
import numpy as np
import torch
import torch.nn.functional as F
from transformers import AutoTokenizer, BertConfig
from Model.MultimodelNER.UMT import UMT
from Model.MultimodelNER import resnet as resnet
from Model.MultimodelNER.resnet_utils import myResnet
from Model.MultimodelNER.VLSP2016.dataset_roberta import convert_mm_examples_to_features, MNERProcessor_2016
from torch.utils.data import (DataLoader, RandomSampler, SequentialSampler,
                              TensorDataset)
from pytorch_pretrained_bert.optimization import BertAdam, warmup_linear
from Model.MultimodelNER.ner_evaluate import evaluate_each_class,evaluate
from seqeval.metrics import classification_report
from tqdm import tqdm, trange
import json
from Model.MultimodelNER.predict import convert_mm_examples_to_features_predict, get_test_examples_predict
from Model.MultimodelNER.Ner_processing import *
CONFIG_NAME = 'bert_config.json'
WEIGHTS_NAME = 'pytorch_model.bin'

logging.basicConfig(format='%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
                    datefmt='%m/%d/%Y %H:%M:%S',
                    level=logging.INFO)
logger = logging.getLogger(__name__)
parser = argparse.ArgumentParser()
## Required parameters
parser.add_argument("--negative_rate",
                    default=16,
                    type=int,
                    help="the negative samples rate")

parser.add_argument('--lamb',
                    default=0.62,
                    type=float)

parser.add_argument('--temp',
                    type=float,
                    default=0.179,
                    help="parameter for CL training")

parser.add_argument('--temp_lamb',
                    type=float,
                    default=0.7,
                    help="parameter for CL training")

parser.add_argument("--data_dir",
                    default='./data/twitter2017',
                    type=str,

                    help="The input data dir. Should contain the .tsv files (or other data files) for the task.")
parser.add_argument("--bert_model", default='vinai/phobert-base-v2', type=str)
parser.add_argument("--task_name",
                    default='sonba',
                    type=str,

                    help="The name of the task to train.")
parser.add_argument("--output_dir",
                    default='Model/MultimodelNER/VLSP2016/best_model/',
                    type=str,
                    help="The output directory where the model predictions and checkpoints will be written.")

## Other parameters
parser.add_argument("--cache_dir",
                    default="",
                    type=str,
                    help="Where do you want to store the pre-trained models downloaded from s3")

parser.add_argument("--max_seq_length",
                    default=128,
                    type=int,
                    help="The maximum total input sequence length after WordPiece tokenization. \n"
                         "Sequences longer than this will be truncated, and sequences shorter \n"
                         "than this will be padded.")

parser.add_argument("--do_train",
                    action='store_true',
                    help="Whether to run training.")

parser.add_argument("--do_eval",
                    action='store_true',
                    help="Whether to run eval on the dev set.")

parser.add_argument("--do_lower_case",
                    action='store_true',
                    help="Set this flag if you are using an uncased model.")

parser.add_argument("--train_batch_size",
                    default=64,
                    type=int,
                    help="Total batch size for training.")

parser.add_argument("--eval_batch_size",
                    default=16,
                    type=int,
                    help="Total batch size for eval.")

parser.add_argument("--learning_rate",
                    default=5e-5,
                    type=float,
                    help="The initial learning rate for Adam.")

parser.add_argument("--num_train_epochs",
                    default=12.0,
                    type=float,
                    help="Total number of training epochs to perform.")

parser.add_argument("--warmup_proportion",
                    default=0.1,
                    type=float,
                    help="Proportion of training to perform linear learning rate warmup for. "
                         "E.g., 0.1 = 10%% of training.")

parser.add_argument("--no_cuda",
                    action='store_true',
                    help="Whether not to use CUDA when available")

parser.add_argument("--local_rank",
                    type=int,
                    default=-1,
                    help="local_rank for distributed training on gpus")

parser.add_argument('--seed',
                    type=int,
                    default=37,
                    help="random seed for initialization")

parser.add_argument('--gradient_accumulation_steps',
                    type=int,
                    default=1,
                    help="Number of updates steps to accumulate before performing a backward/update pass.")

parser.add_argument('--fp16',
                    action='store_true',
                    help="Whether to use 16-bit float precision instead of 32-bit")

parser.add_argument('--loss_scale',
                    type=float, default=0,
                    help="Loss scaling to improve fp16 numeric stability. Only used when fp16 set to True.\n"
                         "0 (default value): dynamic loss scaling.\n"
                         "Positive power of 2: static loss scaling value.\n")

parser.add_argument('--mm_model', default='MTCCMBert', help='model name')  # 'MTCCMBert', 'NMMTCCMBert'
parser.add_argument('--layer_num1', type=int, default=1, help='number of txt2img layer')
parser.add_argument('--layer_num2', type=int, default=1, help='number of img2txt layer')
parser.add_argument('--layer_num3', type=int, default=1, help='number of txt2txt layer')
parser.add_argument('--fine_tune_cnn', action='store_true', help='fine tune pre-trained CNN if True')
parser.add_argument('--resnet_root', default='Model/Resnet/', help='path the pre-trained cnn models')
parser.add_argument('--crop_size', type=int, default=224, help='crop size of image')
parser.add_argument('--path_image', default='Model/MultimodelNER/VLSP2016/Image', help='path to images')
# parser.add_argument('--mm_model', default='TomBert', help='model name') #
parser.add_argument('--server_ip', type=str, default='', help="Can be used for distant debugging.")
parser.add_argument('--server_port', type=str, default='', help="Can be used for distant debugging.")
args = parser.parse_args()



processors = {
    "twitter2015": MNERProcessor_2016,
    "twitter2017": MNERProcessor_2016,
    "sonba": MNERProcessor_2016
}



random.seed(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)


task_name = args.task_name.lower()



processor = processors[task_name]()
label_list = processor.get_labels()
auxlabel_list = processor.get_auxlabels()
num_labels = len(label_list) + 1  # label 0 corresponds to padding, label in label_list starts from 1
auxnum_labels = len(auxlabel_list) + 1  # label 0 corresponds to padding, label in label_list starts from 1

start_label_id = processor.get_start_label_id()
stop_label_id = processor.get_stop_label_id()

# ''' initialization of our conversion matrix, in our implementation, it is a 7*12 matrix initialized as follows:
trans_matrix = np.zeros((auxnum_labels, num_labels), dtype=float)
trans_matrix[0, 0] = 1  # pad to pad
trans_matrix[1, 1] = 1  # O to O
trans_matrix[2, 2] = 0.25  # B to B-MISC
trans_matrix[2, 4] = 0.25  # B to B-PER
trans_matrix[2, 6] = 0.25  # B to B-ORG
trans_matrix[2, 8] = 0.25  # B to B-LOC
trans_matrix[3, 3] = 0.25  # I to I-MISC
trans_matrix[3, 5] = 0.25  # I to I-PER
trans_matrix[3, 7] = 0.25  # I to I-ORG
trans_matrix[3, 9] = 0.25  # I to I-LOC
trans_matrix[4, 10] = 1  # X to X
trans_matrix[5, 11] = 1  # [CLS] to [CLS]
trans_matrix[6, 12] = 1  # [SEP] to [SEP]
'''
trans_matrix = np.zeros((num_labels, auxnum_labels), dtype=float)
trans_matrix[0,0]=1 # pad to pad
trans_matrix[1,1]=1
trans_matrix[2,2]=1
trans_matrix[4,2]=1
trans_matrix[6,2]=1
trans_matrix[8,2]=1
trans_matrix[3,3]=1
trans_matrix[5,3]=1
trans_matrix[7,3]=1
trans_matrix[9,3]=1
trans_matrix[10,4]=1
trans_matrix[11,5]=1
trans_matrix[12,6]=1
'''
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

tokenizer = AutoTokenizer.from_pretrained(args.bert_model, do_lower_case=args.do_lower_case)



net = getattr(resnet, 'resnet152')()
net.load_state_dict(torch.load(os.path.join(args.resnet_root, 'resnet152.pth')))
encoder = myResnet(net, args.fine_tune_cnn, device)


output_model_file = os.path.join(args.output_dir, WEIGHTS_NAME)
# output_config_file = os.path.join(args.output_dir, CONFIG_NAME)
output_encoder_file = os.path.join(args.output_dir, "pytorch_encoder.bin")

temp = args.temp
temp_lamb = args.temp_lamb
lamb = args.lamb
negative_rate = args.negative_rate
# # loadmodel
#     model = UMT.from_pretrained(args.bert_model,
#                                 cache_dir=args.cache_dir, layer_num1=args.layer_num1,
#                                 layer_num2=args.layer_num2,
#                                 layer_num3=args.layer_num3,
#                                 num_labels_=num_labels, auxnum_labels=auxnum_labels)
#     model.load_state_dict(torch.load(output_model_file,map_location=torch.device('cpu')))
#     model.to(device)
#     encoder_state_dict = torch.load(output_encoder_file,map_location=torch.device('cpu'))
#     encoder.load_state_dict(encoder_state_dict)
#     encoder.to(device)
#     print(model)

def load_model(output_model_file, output_encoder_file,encoder,num_labels,auxnum_labels):
    model = UMT.from_pretrained(args.bert_model,
                                cache_dir=args.cache_dir, layer_num1=args.layer_num1,
                                layer_num2=args.layer_num2,
                                layer_num3=args.layer_num3,
                                num_labels_=num_labels, auxnum_labels=auxnum_labels)
    model.load_state_dict(torch.load(output_model_file, map_location=torch.device('cpu')))
    model.to(device)
    encoder_state_dict = torch.load(output_encoder_file, map_location=torch.device('cpu'))
    encoder.load_state_dict(encoder_state_dict)
    encoder.to(device)
    return model, encoder

model_umt,encoder_umt=load_model(output_model_file, output_encoder_file,encoder,num_labels,auxnum_labels)
#
#     # sentence = 'Thương biết_mấy những Thuận, những Liên, những Luận, Xuân, Nghĩa mỗi người một hoàn_cảnh nhưng đều rất giống nhau: rất ham học, rất cố_gắng để đạt mức hiểu biết cao nhất.'
#     # image_path = '/kaggle/working/data/014715.jpg'
#     # # crop_size = 224'
path_image='E:\demo_datn\pythonProject1\Model\MultimodelNER\VLSP2016\Image'
trans_matrix = np.zeros((auxnum_labels,num_labels), dtype=float)
trans_matrix[0,0]=1 # pad to pad
trans_matrix[1,1]=1 # O to O
trans_matrix[2,2]=0.25 # B to B-MISC
trans_matrix[2,4]=0.25 # B to B-PER
trans_matrix[2,6]=0.25 # B to B-ORG
trans_matrix[2,8]=0.25 # B to B-LOC
trans_matrix[3,3]=0.25 # I to I-MISC
trans_matrix[3,5]=0.25 # I to I-PER
trans_matrix[3,7]=0.25 # I to I-ORG
trans_matrix[3,9]=0.25 # I to I-LOC
trans_matrix[4,10]=1   # X to X
trans_matrix[5,11]=1   # [CLS] to [CLS]
trans_matrix[6,12]=1   # [SE
path_image='E:\demo_datn\pythonProject1\Model\MultimodelNER\VLSP2016\Image'

def predict(model_umt, encoder_umt, eval_examples, tokenizer, device,path_image,trans_matrix):

    features = convert_mm_examples_to_features_predict(eval_examples, 256, tokenizer, 224,path_image)

    input_ids = torch.tensor([f.input_ids for f in features], dtype=torch.long)
    input_mask = torch.tensor([f.input_mask for f in features], dtype=torch.long)
    added_input_mask = torch.tensor([f.added_input_mask for f in features], dtype=torch.long)
    segment_ids = torch.tensor([f.segment_ids for f in features], dtype=torch.long)
    img_feats = torch.stack([f.img_feat for f in features])
    print(img_feats)
    eval_data = TensorDataset(input_ids, input_mask, added_input_mask, segment_ids, img_feats)
    eval_sampler = SequentialSampler(eval_data)
    eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=16)

    model_umt.eval()
    encoder_umt.eval()

    y_pred = []
    label_map = {i: label for i, label in enumerate(label_list, 1)}
    label_map[0] = "<pad>"

    for input_ids, input_mask, added_input_mask, segment_ids, img_feats in tqdm(eval_dataloader, desc="Evaluating"):
        input_ids = input_ids.to(device)
        input_mask = input_mask.to(device)
        added_input_mask = added_input_mask.to(device)
        segment_ids = segment_ids.to(device)
        img_feats = img_feats.to(device)

        with torch.no_grad():
            imgs_f, img_mean, img_att = encoder_umt(img_feats)
            predicted_label_seq_ids = model_umt(input_ids, segment_ids, input_mask, added_input_mask, img_att,
                                                trans_matrix)

        logits = predicted_label_seq_ids
        input_mask = input_mask.to('cpu').numpy()

        for i, mask in enumerate(input_mask):
            temp_1 = []
            for j, m in enumerate(mask):
                if j == 0:
                    continue
                if m:
                    if label_map[logits[i][j]] not in ["<pad>", "<s>", "</s>", "X"]:
                        temp_1.append(label_map[logits[i][j]])
                else:
                    break
            y_pred.append(temp_1)

    a = eval_examples[0].text_a.split(" ")

    return y_pred, a

# eval_examples = get_test_examples_predict('E:/demo_datn/pythonProject1/Model/MultimodelNER/VLSP2016/Filetxt/')
# y_pred, a = predict(model_umt, encoder_umt, eval_examples, tokenizer, device,path_image,trans_matrix)
# print(y_pred)
# formatted_output = format_predictions(a, y_pred[0])
# print(formatted_output)
# final= process_predictions(formatted_output)
# final2= combine_entities(final)
# final3= remove_B_prefix(final2)
# final4=combine_i_tags(final3)

# print(final4)