"""
this model has these functions:
1. cluster the training sentences for better mix
2. mix in arbitury layer
3. semi supervised learning 
"""

import argparse
import glob
import logging as log
import os
import random
import time
import torch.nn.functional as F

import numpy as np
import torch
from eval_utils import f1_score, precision_score, recall_score, classification_report, macro_score
from utils import gen_knn_mix_batch
from torch.nn import CrossEntropyLoss
from torch.utils.data import DataLoader, RandomSampler, SequentialSampler, TensorDataset
from torch.utils.data.distributed import DistributedSampler
from tqdm import tqdm, trange
import pickle
from transformers import *
from read_data import *

from tensorboardX import SummaryWriter

from bert_models import BertModel4Mix

logger = log.getLogger(__name__)

use_cuda = torch.cuda.is_available()
#CUDA_VISIBLE_DEVICES=6,7
#os.environ["CUDA_VISIBLE_DEVICES"] = "6,7"
MODEL_CLASSES = {"bert": (BertConfig, BertForTokenClassification, BertTokenizer)}

parser = argparse.ArgumentParser(description='PyTorch BaseNER')
parser.add_argument("--data-dir", default = './data', type = str, required = True)
parser.add_argument("--model-type", default = 'bert', type = str)
parser.add_argument("--model-name", default = 'bert-base-multilingual-cased', type = str)
parser.add_argument("--output-dir", default = './german_eval', type = str)
parser.add_argument('--gpu', default='0,1,2,3', type=str, help='id(s) for CUDA_VISIBLE_DEVICES')
parser.add_argument('--train-examples', default = -1, type = int)

parser.add_argument("--labels", default = "", type = str)
parser.add_argument('--config-name', default = '', type = str)
parser.add_argument("--tokenizer-name", default = '', type = str)
parser.add_argument("--max-seq-length", default = 128, type = int)

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-predict", action="store_true", help="Whether to run predictions on the test set.")
parser.add_argument("--evaluate-during-training", action="store_true", help="Whether to run evaluation during training at each logging step.")
parser.add_argument("--do-lower-case", action="store_true", help="Set this flag if you are using an uncased model.")

parser.add_argument("--batch-size", default = 16, type = int)
parser.add_argument('--eval-batch-size', default = 128, type = int)

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("--learning-rate", default=5e-5, type=float, help="The initial learning rate for Adam.")
parser.add_argument("--weight-decay", default=0.0, type=float, help="Weight decay if we apply some.")
parser.add_argument("--adam-epsilon", default=1e-8, type=float, help="Epsilon for Adam optimizer.")
parser.add_argument("--max-grad-norm", default=1.0, type=float, help="Max gradient norm.")

parser.add_argument("--num-train-epochs", default=20, type=float, help="Total number of training epochs to perform.")
parser.add_argument("--max-steps", default=-1, type=int, help="If > 0: set total number of training steps to perform. Override num_train_epochs.")
parser.add_argument('--warmup-steps', default = 0, type = int,  help="Linear warmup over warmup_steps.")

parser.add_argument('--logging-steps', default = 150, type = int, help="Log every X updates steps.")
parser.add_argument("--save-steps", type=int, default=0, help="Save checkpoint every X updates steps.")
parser.add_argument("--eval-all-checkpoints", action="store_true", help="Evaluate all checkpoints starting with the same prefix as model_name ending and ending with step number")
parser.add_argument("--overwrite-output-dir", action="store_true", help="Overwrite the content of the output directory")

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

parser.add_argument("--pad-subtoken-with-real-label", action="store_true", help="give real label to the padded token instead of `-100` ")
parser.add_argument("--subtoken-label-type",default='real', type=str,help="[real|repeat|O] three ways to do pad subtoken with real label. [real] give the subtoken a real label e.g., B -> B I. [repeat] simply repeat the label e.g., B -> B B. [O] give it a O label. B -> B O")


parser.add_argument("--eval-pad-subtoken-with-first-subtoken-only", action="store_true", help="only works when --pad-subtoken-with-real-label is true, in this mode, we only test the prediction of the first subtoken of each word (if the word could be tokenized into multiple subtoken)")
parser.add_argument("--label-sep-cls", action="store_true", help="label [SEP] [CLS] with special labels, but not [PAD]") 

# inter mix 


# semi
parser.add_argument("--u-batch-size", default = 64, type = int)

parser.add_argument('--T', type = float, default = 1.0,help='sharpen temperature')
parser.add_argument('--sharp', action='store_true')
parser.add_argument('--weight', type = float, default = 1.0)


parser.add_argument("--log-file", default = "results.csv", type = str,help="the file to store resutls")

parser.add_argument("--optimizer", default = "adam", type = str,help='optimizer')
parser.add_argument('--special-label-weight', default=0, type=float, help='the special_label_weight in training . default 0')



# intra-mix




args = parser.parse_args()

os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
use_cuda = torch.cuda.is_available()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
args.device = device
args.n_gpu = torch.cuda.device_count()
print("gpu num: ", args.n_gpu)

best_f1 = 0





def set_seed(args):
    logger.info("random seed %s", args.seed)
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    if len(args.gpu) > 0:
        torch.cuda.manual_seed_all(args.seed)


    
def linear_rampup(current, rampup_length=args.num_train_epochs):
    if rampup_length == 0:
        return 1.0
    else:
        current = np.clip(current / rampup_length, 0.0, 1.0)
        return float(current)

def read_examples_from_ontonotes_file(args, tokenizer, labelslist, pad_token_label_id, mode, 
              omit_sep_cls_token=False,
              pad_subtoken_with_real_label=False):
    file_path = os.path.join(args.data_dir, 'ontotest.ner')
    guid_index = 0
    examples = []
    
    label_map = {"PER": "PER",
             "person": "PER",
             "PERSON": "PER",
             "DOCTOR": "PER",
             "PATIENT": "PER",
             "LOC": "LOC",
             "location": "LOC",
             "GPE": "LOC",
             "FAC": "LOC",
             "HOSPITAL": "LOC", 
             "CITY": "LOC",
             "STATE": "LOC",
             "COUNTRY": "LOC",
             "LOCATION_OTHER": "LOC",
             "ORG": "ORG",
             "corporation": "ORG",
             "group": "ORG",
             "ORGANIZATION": "ORG",
             "MISC": "MISC",
             "product": "MISC",
             "creative-work": "MISC",
             "PRODUCT": "MISC", 
             "NORP": "MISC",
             "EVENT": "MISC",
             "LANGUAGE": "MISC",
             "LAW": "MISC",
             "WORK_OF_ART": "MISC"}
        
    
    with open(file_path, encoding="utf-8") as f:
        words = []
        labels = []
        for line in f:
            if line.startswith("-DOCSTART-") or line == "" or line == "\n":
                if words:
                    examples.append(InputExample(
                        guid="{}-{}".format(mode, guid_index), words=words, labels=labels))
                    guid_index += 1
                    words = []
                    labels = []
                    
            else:
                elems = line.strip().split()
                label = elems[-1].replace("\n", "")
                words.append(elems[0])
               
                if label != "O":
                    iob_prefix = label[:2]
                    etype = label[2:]
                    if etype in label_map:
                        labels.append(iob_prefix + label_map[etype])
                    else:
                        labels.append("O")
                        
                else:
                    labels.append("O")
        if words:
            examples.append(InputExample(
                guid="{}-{}".format(mode, guid_index), words=words, labels=labels))
    entityheavyexamples=[]
    guid_index=0
    for e in examples:
        count=0
        for l in e.labels:
            if l!="O":
                count=count+1
            
        if count/len( e.labels)>0.5:
            e.guid="{}-{}".format(mode, guid_index)
            entityheavyexamples.append(e)
            guid_index += 1
            if len(entityheavyexamples)>49:
                break

        print(e.guid,"   ", e.words)
        print(e.labels)
    
        #examples = examples[0]
        print(mode)
        print('data num: {}'.format(len(examples)))

    features = convert_examples_to_features(entityheavyexamples, labelslist, args.max_seq_length, tokenizer, 
                                cls_token = tokenizer.cls_token, sep_token =  tokenizer.sep_token, pad_token = tokenizer.convert_tokens_to_ids([tokenizer.pad_token])[0],
                                cls_token_segment_id = 2 if args.model_type in ["xlnet"] else 0, 
                                sequence_a_segment_id = 0, pad_token_segment_id=4 if args.model_type in ["xlnet"] else 0, 
                                pad_token_label_id = pad_token_label_id,
                                omit_sep_cls_token=omit_sep_cls_token,
                                pad_subtoken_with_real_label=pad_subtoken_with_real_label,
                                subtoken_label_type=args.subtoken_label_type,
                                label_sep_cls=args.label_sep_cls)
    
    all_input_ids = torch.tensor([f.input_ids for f in features], dtype=torch.long)
    all_input_mask = torch.tensor([f.input_mask for f in features], dtype=torch.long)
    all_segment_ids = torch.tensor([f.segment_ids for f in features], dtype=torch.long)
    all_label_ids = torch.tensor([f.label_ids for f in features], dtype=torch.long)
    all_subtoken_ids = torch.tensor([f.subtoken_ids for f in features], dtype=torch.long)
    all_sent_id = torch.tensor([f.sent_id for f in features], dtype=torch.long)
    dataset = TensorDataset(all_input_ids, all_input_mask, all_segment_ids, all_label_ids, all_subtoken_ids,all_sent_id)

    
    return dataset


def train(args,train_dataset, eval_dataset, test_dataset_regular,test_dataset_challenging, model, tokenizer, labels, pad_token_label_id):
    
    global best_f1
    tb_writer = SummaryWriter()
    print('tb_writer.logdir',tb_writer.logdir)
    
    train_dataloader = DataLoader(train_dataset, batch_size = args.batch_size, shuffle = True)
    labeled_dataloader = train_dataloader


    if args.max_steps > 0:
        t_total = args.max_steps
        args.num_train_epochs = args.max_steps // (len(train_dataloader) // args.gradient_accumulation_steps) + 1
    else:
        t_total = len(train_dataloader) // args.gradient_accumulation_steps * args.num_train_epochs

    no_decay = ["bias", "LayerNorm.weight"]
    optimizer_grouped_parameters = [
        {
            "params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
            "weight_decay": args.weight_decay,
        },
        {"params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], "weight_decay": 0.0},
    ]


    if args.n_gpu > 1:
        model = torch.nn.DataParallel(model)
    if args.optimizer=='adam':
        optimizer = AdamW(optimizer_grouped_parameters, lr=args.learning_rate, eps=args.adam_epsilon)
    elif args.optimizer=='sgd':
        optimizer = torch.optim.SGD(model.parameters(), lr=args.learning_rate)
        
    scheduler = get_linear_schedule_with_warmup(
        optimizer, num_warmup_steps=args.warmup_steps, num_training_steps=t_total
    )    
    logger.info("***** Running training *****")
    logger.info("  Num examples = %d", len(train_dataset))
    logger.info("  Num Epochs = %d", args.num_train_epochs)
    logger.info(
        "  Total train batch size (w. parallel, accumulation) = %d",
        args.batch_size
        * args.gradient_accumulation_steps),
    logger.info("  Gradient Accumulation steps = %d", args.gradient_accumulation_steps)
    logger.info("  Total optimization steps = %d", t_total)

    global_step = 0
    epochs_trained = 0
    steps_trained_in_current_epoch = 0

    tr_loss, logging_loss = 0.0, 0.0

    
    #eval_f1 = []
    test_f1 = []
    test_f1_regular = []
    test_f1_challenging = []
    model.zero_grad()

    train_iterator = trange(epochs_trained, int(args.num_train_epochs), desc='Epoch')
    set_seed(args)



    for epoch in train_iterator:

        #Make sure batch tuple indeces match ones in train
        print("Epoch #",epoch)
        epoch_iterator = tqdm(train_dataloader, desc="Iteration")        
        for step, batch in enumerate(epoch_iterator):
            if steps_trained_in_current_epoch > 0:
                steps_trained_in_current_epoch -= 1
                continue
            model.train()
            
            batch = tuple(t.to(args.device) for t in batch)

            #inputs_c = {"input_ids": batch[0].to(args.device),"attention_mask": batch[1].to(args.device),'subtoken_ids':batch[4].to(args.device)}
            #target_c=batch[3].to(args.device) 

            b_input_ids = batch[0].to(device)
            b_input_mask = batch[1].to(device)
            b_labels = batch[3].to(device)         
                                    
            model.zero_grad()
          
            result = model(b_input_ids, 
                        token_type_ids=None, 
                        attention_mask=b_input_mask, 
                        labels=b_labels)
            
            loss = result[0]
           
            if args.gradient_accumulation_steps > 1:
                loss = loss / args.gradient_accumulation_steps
            
            loss.backward()
            tr_loss += loss.item()

            if (step + 1) % args.gradient_accumulation_steps == 0:
                torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)
                optimizer.step()
                scheduler.step()  
                model.zero_grad()
                global_step += 1

                if args.logging_steps > 0 and global_step % args.logging_steps == 0:
                    # Log metrics
                    if (args.evaluate_during_training):
                        
                        results, _ = evaluate(args, model, tokenizer, labels, pad_token_label_id, eval_dataset, parallel = False, mode="dev", prefix = str(global_step))
                        for key, value in results.items():
                            tb_writer.add_scalar("eval_{}".format(key), value, global_step)
                        logger.info("Model name: %s", args.output_dir)
                        logger.info("Epoch is %s", epoch)
                        if results['f1'] >= best_f1:
                            best_f1 = results['f1']
                            results, _ = evaluate(args, model, tokenizer, labels, pad_token_label_id, test_dataset_regular, parallel = False, mode="test", prefix = str(global_step))
                            test_f1_regular.append(results['f1'])
                            results, _ = evaluate(args,model, tokenizer, labels, pad_token_label_id, test_dataset_challenging, parallel = False, mode="test", prefix = str(global_step))
                            test_f1_challenging.append(results['f1'])
                            
                            
                            output_dir = os.path.join(args.output_dir, "best")
                            #epochinfo=os.path.join(output_dir, "bestepochinfo.txt")
                            #epochoutput = open(epochinfo, 'a')
                            #print("Best Epoch #",epoch, file = epochoutput)
                            #epochoutput.close()
                            if not os.path.exists(output_dir):
                                os.makedirs(output_dir)
                            logger.info("Saving best model to %s", output_dir)
                            logger.info("Epochs trained is %s", epochs_trained)
                            logger.info("Epoch is %s", epoch)
                            model_to_save = (
                                model.module if hasattr(model, "module") else model)  
                            model_to_save.save_pretrained(output_dir)
                            tokenizer.save_pretrained(output_dir)
                            torch.save(args, os.path.join(output_dir, "training_args.bin"))
                            logger.info("Saving model checkpoint to %s", output_dir)

                            torch.save(optimizer.state_dict(), os.path.join(output_dir, "optimizer.pt"))
                            torch.save(scheduler.state_dict(), os.path.join(output_dir, "scheduler.pt"))
                            logger.info("Saving optimizer and scheduler states to %s", output_dir)
                    
                    tb_writer.add_scalar("lr", scheduler.get_last_lr()[0], global_step)
                    tb_writer.add_scalar("loss", (tr_loss - logging_loss) / args.logging_steps, global_step)
                    logging_loss = tr_loss

                    logger.info("logging train info!!!")
                    logger.info("*")

            
        # eval and save the best model based on dev set after each epoch
        if (args.evaluate_during_training):
            
            results, _ = evaluate(args, model, tokenizer, labels, pad_token_label_id, eval_dataset, parallel = False, mode="dev", prefix = str(global_step))
            for key, value in results.items():
                tb_writer.add_scalar("eval_{}".format(key), value, global_step)

            
            if results['f1'] >= best_f1:
                best_f1 = results['f1']
                results, _ = evaluate(args, model, tokenizer, labels, pad_token_label_id, test_dataset_regular, parallel = False, mode="test", prefix = str(global_step))
                test_f1_regular.append(results['f1'])
                results, _ = evaluate(args,model, tokenizer, labels, pad_token_label_id, test_dataset_challenging, parallel = False, mode="test", prefix = str(global_step))
                test_f1_challenging.append(results['f1'])
                
                output_dir = os.path.join(args.output_dir, "best")
                if not os.path.exists(output_dir):
                    os.makedirs(output_dir)
                logger.info("Epoch #", epoch)
                logger.info("Saving best model to %s", output_dir)
                logger.info("Epochs trained is %s", epochs_trained)
                logger.info("Epoch is %s", epoch)
                model_to_save = (model.module if hasattr(model, "module") else model)  
                model_to_save.save_pretrained(output_dir)
                tokenizer.save_pretrained(output_dir)
                torch.save(args, os.path.join(output_dir, "training_args.bin"))
                logger.info("Saving model checkpoint to %s", output_dir)

                torch.save(optimizer.state_dict(), os.path.join(output_dir, "optimizer.pt"))
                torch.save(scheduler.state_dict(), os.path.join(output_dir, "scheduler.pt"))
                logger.info("Saving optimizer and scheduler states to %s", output_dir)

        if args.max_steps > 0 and global_step > args.max_steps:
            train_iterator.close()
            break
        logger.info("Epoch is %s", epoch)
    args.tb_writer_logdir=tb_writer.logdir
    tb_writer.close()
    return global_step, tr_loss / global_step, test_f1_regular ,test_f1_challenging

def output_eval_results(out_label_list,preds_list,input_id_list,file_name):
    with open(file_name,'w') as fout:
        for i in range(len(out_label_list)):
            label=out_label_list[i]
            pred=preds_list[i]
            tokens=input_id_list[i]
            for j in range(len(label)):
                if tokens[j]=='[PAD]':
                    continue
                fout.write('{}\t{}\t{}\n'.format(tokens[j] ,label[j],pred[j]))
            fout.write('\n')


def evaluate(args,model, tokenizer, labels, pad_token_label_id,  eval_dataset = None, parallel = True, mode = 'dev', prefix = ''):
    if eval_dataset is None and mode=='dev':
        eval_dataset = read_data(args, tokenizer, labels, pad_token_label_id, mode = mode)

    eval_dataloader = DataLoader(eval_dataset, batch_size = args.eval_batch_size, shuffle = False)

    if parallel:
        model = torch.nn.DataParallel(model)

    # Eval!
    logger.info("***** Running evaluation %s *****", mode + '-' + prefix)
    logger.info("  Num examples = %d", len(eval_dataset))
    logger.info("  Batch size = %d", args.eval_batch_size)

    eval_loss = 0.0
    nb_eval_steps = 0
    preds = None
    out_label_ids = None

    all_subtoken_ids=None
    model.eval()

    for batch in tqdm(eval_dataloader, desc="Evaluating"):
        batch = tuple(t.to(args.device) for t in batch)

        with torch.no_grad():
            inputs = {"input_ids": batch[0], "attention_mask": batch[1], "labels": batch[3],'sent_id' : batch[5]}
            inputs["token_type_ids"] = batch[2]
            target=inputs['labels']
            
            result  = model(inputs['input_ids'], 
                            token_type_ids=None, 
                            attention_mask=inputs["attention_mask"])
     
            logits=result[0]

          
        nb_eval_steps += 1
        
        if preds is None:
            preds = logits.detach().cpu().numpy()
            out_label_ids = inputs["labels"].detach().cpu().numpy()
            all_subtoken_ids=batch[4].detach().cpu().numpy()
            sent_id=inputs['sent_id'].detach().cpu().numpy()
            input_ids=inputs['input_ids'].detach().cpu().numpy()
        else:
       
            preds = np.append(preds, logits.detach().cpu().numpy(), axis=0)
            out_label_ids = np.append(out_label_ids, inputs["labels"].detach().cpu().numpy(), axis=0)
            all_subtoken_ids = np.append(all_subtoken_ids, batch[4].detach().cpu().numpy(), axis=0)
            sent_id = np.append(sent_id, inputs['sent_id'].detach().cpu().numpy(), axis=0)
            input_ids= np.append(input_ids, inputs['input_ids'].detach().cpu().numpy(), axis=0)

    #eval_loss = eval_loss / nb_eval_steps
    preds = np.argmax(preds, axis=2)

    label_map = {i: label for i, label in enumerate(labels)}

    out_label_list = [[] for _ in range(out_label_ids.shape[0])]
    preds_list = [[] for _ in range(out_label_ids.shape[0])]
    input_id_list = [[] for _ in range(input_ids.shape[0])]

    for i in range(out_label_ids.shape[0]):
        for j in range(out_label_ids.shape[1]):
            if args.pad_subtoken_with_real_label  or args.label_sep_cls:

                if args.eval_pad_subtoken_with_first_subtoken_only:
                    if all_subtoken_ids[i,j] ==1: 
                        out_label_list[i].append(label_map[out_label_ids[i][j]])
                        preds_list[i].append(label_map[preds[i][j]])
                        tid=input_ids[i][j]
                        input_id_list[i].append(tokenizer.convert_ids_to_tokens([tid])[0])


                else:
                    if all_subtoken_ids[i,j] in [0,1] and out_label_ids[i, j] != pad_token_label_id:# in this case, we consider all the tokens.
                        out_label_list[i].append(label_map[out_label_ids[i][j]])
                        preds_list[i].append(label_map[preds[i][j]])            
                        input_id_list[i].append(tokenizer.convert_ids_to_tokens([input_ids[i][j]]))    
            else:
                if all_subtoken_ids[i,j] in [0,1] and out_label_ids[i, j] != pad_token_label_id:# in this case, we consider all the tokens.
                    out_label_list[i].append(label_map[out_label_ids[i][j]])
                    preds_list[i].append(label_map[preds[i][j]])                
                    input_id_list[i].append(tokenizer.convert_ids_to_tokens([input_ids[i][j]]))
    file_name=os.path.join(args.output_dir,'{}_pred_results.tsv'.format(mode))
    output_eval_results(out_label_list,preds_list,input_id_list,file_name)

    macro_scores=macro_score(out_label_list, preds_list)
    results = {
      
        "precision": precision_score(out_label_list, preds_list),
        "recall": recall_score(out_label_list, preds_list),
        "f1": f1_score(out_label_list, preds_list),
        'macro_f1':macro_scores['macro_f1'],
        'macro_precision':macro_scores['macro_precision'],
        'macro_recall':macro_scores['macro_recall']
    }

    logger.info("***** Eval results %s *****", mode + '-' + prefix)

    for key in sorted(results.keys()):
        logger.info("  %s = %s", key, str(results[key]))

    return results, preds_list

def main():
    global best_f1
    if (os.path.exists(args.output_dir) and os.listdir(args.output_dir) and args.do_train and not args.overwrite_output_dir):
        raise ValueError( "Output directory ({}) already exists and is not empty. Use --overwrite_output_dir to overcome.".format(args.output_dir))
    
    logger.setLevel(log.INFO)
    formatter = log.Formatter("%(asctime)s - %(levelname)s - %(name)s -   %(message)s", datefmt="%m/%d/%Y %H:%M:%S")
    
    if not os.path.exists(args.output_dir):
            os.makedirs(args.output_dir)
            
            
    fh = log.FileHandler(args.output_dir  +'/' + str(args.train_examples)+'-' + 'log.txt')
    fh.setLevel(log.INFO)
    fh.setFormatter(formatter)

    ch = log.StreamHandler()
    ch.setLevel(log.INFO)
    ch.setFormatter(formatter)

    logger.addHandler(ch)
    logger.addHandler(fh)
    
    logger.info("------NEW RUN-----")

    logger.info("device: %s, n_gpu: %s", args.device, args.n_gpu)

    set_seed(args)

    labels = get_labels(args.labels)
    num_labels = len(labels)
    args.num_labels=num_labels

    pad_token_label_id = CrossEntropyLoss().ignore_index





    args.model_type = args.model_type.lower()
    config_class, model_class, tokenizer_class = MODEL_CLASSES[args.model_type]

    
    tokenizer = tokenizer_class.from_pretrained(args.model_name, do_lower_case=args.do_lower_case)

    
    config = config_class.from_pretrained(
        args.config_name if args.config_name else args.model_name,
        num_labels=num_labels,
    )
    tokenizer = tokenizer_class.from_pretrained(
        args.tokenizer_name if args.tokenizer_name else args.model_name,
        do_lower_case=args.do_lower_case,
    )
    model = BertForTokenClassification.from_pretrained(
        args.model_name,
        from_tf=bool(".ckpt" in args.model_name),
        config=config)
   
    model.to(args.device)

    logger.info("Training/evaluation parameters %s", args)
              
    
    if args.do_predict:
        print("Doing Predict!!!!")
        test_dataset_ontonotes = read_examples_from_ontonotes_file(args, tokenizer, labels, pad_token_label_id, mode = 'test', pad_subtoken_with_real_label=args.pad_subtoken_with_real_label)
    
        output_dir = os.path.join(args.output_dir, "best")
        tokenizer = tokenizer_class.from_pretrained(output_dir, do_lower_case=args.do_lower_case)
        

        model = model_class.from_pretrained(output_dir)
        model.to(args.device)
        result, predictions = evaluate(args, model, tokenizer, labels, pad_token_label_id,test_dataset_ontonotes, mode="test", prefix = 'final')
        print("Ontonotes Test set Results: ",result)
  




main()

    