import logging
logging.basicConfig(
    format='%(asctime)s - %(levelname)s - %(name)s -  %(message)s',
    datefmt='%Y/%m/%d %H:%M:%S',
    level=logging.INFO,
    )
logger = logging.getLogger("Main")

import os,random
import numpy as np
import torch
from processing import convert_examples_to_features, read_squad_examples
from processing import ChineseFullTokenizer
from pytorch_pretrained_bert.my_modeling import BertConfig
from optimization import BERTAdam
import config
from utils import read_and_convert, divide_parameters
from modeling import BertForQASimple, BertForQASimpleAdaptorTraining
from textbrewer import DistillationConfig, TrainingConfig, BasicTrainer
from torch.utils.data import TensorDataset, DataLoader, RandomSampler
from functools import partial

from train_eval import predict

def args_check(args):
    if os.path.exists(args.output_dir) and os.listdir(args.output_dir):
        logger.warning("Output directory () already exists and is not empty.")
    if args.gradient_accumulation_steps < 1:
        raise ValueError("Invalid gradient_accumulation_steps parameter: {}, should be >= 1".format(
                            args.gradient_accumulation_steps))

    if not args.do_train and not args.do_predict:
        raise ValueError("At least one of `do_train` or `do_predict` must be True.")

    if args.local_rank == -1 or args.no_cuda:
        device = torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu")
        n_gpu = torch.cuda.device_count() if not args.no_cuda else 0
    else:
        device = torch.device("cuda", args.local_rank)
        n_gpu = 1
        torch.distributed.init_process_group(backend='nccl')
    logger.info("device %s n_gpu %d distributed training %r", device, n_gpu, bool(args.local_rank != -1))
    args.n_gpu = n_gpu
    args.device = device
    return device, n_gpu

def main():
    #parse arguments
    config.parse()
    args = config.args
    for k,v in vars(args).items():
        logger.info(f"{k}:{v}")
    #set seeds
    torch.manual_seed(args.random_seed)
    torch.cuda.manual_seed_all(args.random_seed)
    np.random.seed(args.random_seed)
    random.seed(args.random_seed)

    #arguments check
    device, n_gpu = args_check(args)
    os.makedirs(args.output_dir, exist_ok=True)
    forward_batch_size = int(args.train_batch_size / args.gradient_accumulation_steps)
    args.forward_batch_size = forward_batch_size

    #load bert config
    bert_config_S = BertConfig.from_json_file(args.bert_config_file_S)
    assert args.max_seq_length <= bert_config_S.max_position_embeddings

    #read data
    train_examples = None
    train_features = None
    eval_examples = None
    eval_features = None
    num_train_steps = None

    tokenizer = ChineseFullTokenizer(vocab_file=args.vocab_file, do_lower_case=args.do_lower_case)
    convert_fn = partial(convert_examples_to_features,
                         tokenizer=tokenizer,
                         max_seq_length=args.max_seq_length,
                         doc_stride=args.doc_stride,
                         max_query_length=args.max_query_length)
    if args.do_train:
        train_examples,train_features = read_and_convert(args.train_file,is_training=True, do_lower_case=args.do_lower_case,
                                                         read_fn=read_squad_examples,convert_fn=convert_fn)
        if args.fake_file_1:
            fake_examples1,fake_features1 = read_and_convert(args.fake_file_1,is_training=True, do_lower_case=args.do_lower_case,
                                                             read_fn=read_squad_examples,convert_fn=convert_fn)
            train_examples += fake_examples1
            train_features += fake_features1
        if args.fake_file_2:
            fake_examples2, fake_features2 = read_and_convert(args.fake_file_2,is_training=True, do_lower_case=args.do_lower_case,
                                                              read_fn=read_squad_examples,convert_fn=convert_fn)
            train_examples += fake_examples2
            train_features += fake_features2

        num_train_steps = int(len(train_features)/args.train_batch_size) * args.num_train_epochs

    if args.do_predict:
        eval_examples,eval_features = read_and_convert(args.predict_file,is_training=False, do_lower_case=args.do_lower_case,
                                                         read_fn=read_squad_examples,convert_fn=convert_fn)

    #Build Model and load checkpoint
    model_S = BertForQASimple(bert_config_S,args)
    #Load student
    if args.load_model_type=='bert':
        assert args.init_checkpoint_S is not None
        state_dict_S = torch.load(args.init_checkpoint_S, map_location='cpu')
        state_weight = {k[5:]:v for k,v in state_dict_S.items() if k.startswith('bert.')}
        missing_keys,_ = model_S.bert.load_state_dict(state_weight,strict=False)
        assert len(missing_keys)==0
    elif args.load_model_type=='all':
        assert args.tuned_checkpoint_S is not None
        state_dict_S = torch.load(args.tuned_checkpoint_S,map_location='cpu')
        model_S.load_state_dict(state_dict_S)
    else:
        logger.info("Model is randomly initialized.")
    model_S.to(device)

    if args.local_rank != -1 or n_gpu > 1:
        if args.local_rank != -1:
            raise NotImplementedError
        elif n_gpu > 1:
            model_S = torch.nn.DataParallel(model_S) #,output_device=n_gpu-1)

    if args.do_train:
        #parameters
        params = list(model_S.named_parameters())
        all_trainable_params = divide_parameters(params, lr=args.learning_rate)
        logger.info("Length of all_trainable_params: %d", len(all_trainable_params))


        optimizer = BERTAdam(all_trainable_params,lr=args.learning_rate,
                             warmup=args.warmup_proportion,t_total=num_train_steps,schedule=args.schedule,
                             s_opt1=args.s_opt1, s_opt2=args.s_opt2, s_opt3=args.s_opt3)

        logger.info("***** Running training *****")
        logger.info("  Num orig examples = %d", len(train_examples))
        logger.info("  Num split examples = %d", len(train_features))
        logger.info("  Forward batch size = %d", forward_batch_size)
        logger.info("  Num backward steps = %d", num_train_steps)

        ########### DISTILLATION ###########
        train_config = TrainingConfig(
            gradient_accumulation_steps = args.gradient_accumulation_steps,
            ckpt_frequency = args.ckpt_frequency,
            log_dir = args.output_dir,
            output_dir = args.output_dir,
            device = args.device)

        distiller = BasicTrainer(train_config = train_config,
                                   model = model_S,
                                   adaptor = BertForQASimpleAdaptorTraining)

        all_input_ids = torch.tensor([f.input_ids for f in train_features], dtype=torch.long)
        all_input_mask = torch.tensor([f.input_mask for f in train_features], dtype=torch.long)
        all_doc_mask = torch.tensor([f.doc_mask for f in train_features], dtype=torch.float)
        all_segment_ids = torch.tensor([f.segment_ids for f in train_features], dtype=torch.long)
        all_start_positions = torch.tensor([f.start_position for f in train_features], dtype=torch.long)
        all_end_positions = torch.tensor([f.end_position for f in train_features], dtype=torch.long)

        train_dataset = TensorDataset(all_input_ids, all_segment_ids, all_input_mask, all_doc_mask,
                                   all_start_positions, all_end_positions)
        if args.local_rank == -1:
            train_sampler = RandomSampler(train_dataset)
        else:
            raise NotImplementedError
        train_dataloader = DataLoader(train_dataset, sampler=train_sampler, batch_size=args.forward_batch_size,drop_last=True)
        callback_func = partial(predict, 
                eval_examples=eval_examples,
                eval_features=eval_features,
                args=args)
        with distiller:
            distiller.train(optimizer, scheduler=None, dataloader=train_dataloader,
                              num_epochs=args.num_train_epochs, callback=callback_func)

    if not args.do_train and args.do_predict:
        res = predict(model_S,eval_examples,eval_features,step=0,args=args)
        print (res)




if __name__ == "__main__":
    main()

[((2, 0, 6, 5), 'logging.basicConfig', 'logging.basicConfig', (), '', False, 'import logging\n'), ((7, 9, 7, 34), 'logging.getLogger', 'logging.getLogger', ({(7, 27, 7, 33): '"""Main"""'}, {}), "('Main')", False, 'import logging\n'), ((49, 4, 49, 18), 'config.parse', 'config.parse', ({}, {}), '()', False, 'import config\n'), ((54, 4, 54, 39), 'torch.manual_seed', 'torch.manual_seed', ({(54, 22, 54, 38): 'args.random_seed'}, {}), '(args.random_seed)', False, 'import torch\n'), ((55, 4, 55, 48), 'torch.cuda.manual_seed_all', 'torch.cuda.manual_seed_all', ({(55, 31, 55, 47): 'args.random_seed'}, {}), '(args.random_seed)', False, 'import torch\n'), ((56, 4, 56, 36), 'numpy.random.seed', 'np.random.seed', ({(56, 19, 56, 35): 'args.random_seed'}, {}), '(args.random_seed)', True, 'import numpy as np\n'), ((57, 4, 57, 33), 'random.seed', 'random.seed', ({(57, 16, 57, 32): 'args.random_seed'}, {}), '(args.random_seed)', False, 'import os, random\n'), ((61, 4, 61, 47), 'os.makedirs', 'os.makedirs', (), '', False, 'import os, random\n'), ((66, 20, 66, 70), 'pytorch_pretrained_bert.my_modeling.BertConfig.from_json_file', 'BertConfig.from_json_file', ({(66, 46, 66, 69): 'args.bert_config_file_S'}, {}), '(args.bert_config_file_S)', False, 'from pytorch_pretrained_bert.my_modeling import BertConfig\n'), ((76, 16, 76, 98), 'processing.ChineseFullTokenizer', 'ChineseFullTokenizer', (), '', False, 'from processing import ChineseFullTokenizer\n'), ((77, 17, 81, 64), 'functools.partial', 'partial', (), '', False, 'from functools import partial\n'), ((103, 14, 103, 49), 'modeling.BertForQASimple', 'BertForQASimple', ({(103, 30, 103, 43): 'bert_config_S', (103, 44, 103, 48): 'args'}, {}), '(bert_config_S, args)', False, 'from modeling import BertForQASimple, BertForQASimpleAdaptorTraining\n'), ((26, 7, 26, 38), 'os.path.exists', 'os.path.exists', ({(26, 22, 26, 37): 'args.output_dir'}, {}), '(args.output_dir)', False, 'import os, random\n'), ((26, 43, 26, 70), 'os.listdir', 'os.listdir', ({(26, 54, 26, 69): 'args.output_dir'}, {}), '(args.output_dir)', False, 'import os, random\n'), ((39, 17, 39, 54), 'torch.device', 'torch.device', ({(39, 30, 39, 36): '"""cuda"""', (39, 38, 39, 53): 'args.local_rank'}, {}), "('cuda', args.local_rank)", False, 'import torch\n'), ((41, 8, 41, 60), 'torch.distributed.init_process_group', 'torch.distributed.init_process_group', (), '', False, 'import torch\n'), ((83, 40, 84, 107), 'utils.read_and_convert', 'read_and_convert', (), '', False, 'from utils import read_and_convert, divide_parameters\n'), ((99, 38, 100, 107), 'utils.read_and_convert', 'read_and_convert', (), '', False, 'from utils import read_and_convert, divide_parameters\n'), ((107, 23, 107, 77), 'torch.load', 'torch.load', (), '', False, 'import torch\n'), ((128, 31, 128, 79), 'utils.divide_parameters', 'divide_parameters', (), '', False, 'from utils import read_and_convert, divide_parameters\n'), ((132, 20, 134, 88), 'optimization.BERTAdam', 'BERTAdam', (), '', False, 'from optimization import BERTAdam\n'), ((143, 23, 148, 33), 'textbrewer.TrainingConfig', 'TrainingConfig', (), '', False, 'from textbrewer import DistillationConfig, TrainingConfig, BasicTrainer\n'), ((150, 20, 152, 76), 'textbrewer.BasicTrainer', 'BasicTrainer', (), '', False, 'from textbrewer import DistillationConfig, TrainingConfig, BasicTrainer\n'), ((154, 24, 154, 93), 'torch.tensor', 'torch.tensor', (), '', False, 'import torch\n'), ((155, 25, 155, 95), 'torch.tensor', 'torch.tensor', (), '', False, 'import torch\n'), ((156, 23, 156, 92), 'torch.tensor', 'torch.tensor', (), '', False, 'import torch\n'), ((157, 26, 157, 97), 'torch.tensor', 'torch.tensor', (), '', False, 'import torch\n'), ((158, 30, 158, 104), 'torch.tensor', 'torch.tensor', (), '', False, 'import torch\n'), ((159, 28, 159, 100), 'torch.tensor', 'torch.tensor', (), '', False, 'import torch\n'), ((161, 24, 162, 74), 'torch.utils.data.TensorDataset', 'TensorDataset', ({(161, 38, 161, 51): 'all_input_ids', (161, 53, 161, 68): 'all_segment_ids', (161, 70, 161, 84): 'all_input_mask', (161, 86, 161, 98): 'all_doc_mask', (162, 35, 162, 54): 'all_start_positions', (162, 56, 162, 73): 'all_end_positions'}, {}), '(all_input_ids, all_segment_ids, all_input_mask, all_doc_mask,\n    all_start_positions, all_end_positions)', False, 'from torch.utils.data import TensorDataset, DataLoader, RandomSampler\n'), ((167, 27, 167, 126), 'torch.utils.data.DataLoader', 'DataLoader', (), '', False, 'from torch.utils.data import TensorDataset, DataLoader, RandomSampler\n'), ((168, 24, 171, 26), 'functools.partial', 'partial', (), '', False, 'from functools import partial\n'), ((177, 14, 177, 75), 'train_eval.predict', 'predict', (), '', False, 'from train_eval import predict\n'), ((37, 16, 37, 41), 'torch.cuda.device_count', 'torch.cuda.device_count', ({}, {}), '()', False, 'import torch\n'), ((86, 44, 87, 111), 'utils.read_and_convert', 'read_and_convert', (), '', False, 'from utils import read_and_convert, divide_parameters\n'), ((91, 45, 92, 112), 'utils.read_and_convert', 'read_and_convert', (), '', False, 'from utils import read_and_convert, divide_parameters\n'), ((113, 23, 113, 77), 'torch.load', 'torch.load', (), '', False, 'import torch\n'), ((164, 28, 164, 56), 'torch.utils.data.RandomSampler', 'RandomSampler', ({(164, 42, 164, 55): 'train_dataset'}, {}), '(train_dataset)', False, 'from torch.utils.data import TensorDataset, DataLoader, RandomSampler\n'), ((123, 22, 123, 52), 'torch.nn.DataParallel', 'torch.nn.DataParallel', ({(123, 44, 123, 51): 'model_S'}, {}), '(model_S)', False, 'import torch\n'), ((36, 40, 36, 65), 'torch.cuda.is_available', 'torch.cuda.is_available', ({}, {}), '()', False, 'import torch\n')]