# coding=utf-8
# Copyright 2019 The HuggingFace Inc. team.
# Copyright (c) 2019 The HuggingFace Inc.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
""" Finetuning seq2seq models for sequence generation."""

import argparse
import functools
import logging
import os
import random
import sys
sys.path.append(r'../')


import numpy as np
from tqdm import tqdm, trange
import torch
from torch.optim import Adam
from torch.utils.data import DataLoader, RandomSampler, SequentialSampler

from transformers import (
    AutoTokenizer,
    BertForMaskedLM,
    BertConfig,
    PreTrainedEncoderDecoder,
    Model2Models,
)

from utils_summarization import (
    CNNDailyMailDataset,
    encode_for_summarization,
    fit_to_block_size,
    build_lm_labels,
    build_mask,
    compute_token_type_ids,
)

from utils_chemistry import (ChemistryDataset,)
'''
class InputExample(object):
    def __init__(self,example_id,question_input,question_varible_output=None,condition_output=None):
        self.example_id=example_id
        self.question_input=question_input
        self.question_varible_output=question_varible_output
        self.condition_output=condition_output
'''


logger = logging.getLogger(__name__)
logging.basicConfig(stream=sys.stdout, level=logging.INFO)


def set_seed(args):
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)


# ------------
# Load dataset
# ------------


def load_and_cache_examples(args, tokenizer, prefix="train"):
    dataset = ChemistryDataset(tokenizer, prefix=prefix, data_dir=args.data_dir)
    return dataset


def collate(data, tokenizer, input_block_size,output_block_size):
    """ List of tuple as an input. """
    question_inputs=[]
    question_varible_outputs=[]
    condition_outputs=[]
    for i,example in enumerate(data):
        question_input=tokenizer.encode(example.question_input)
        question_input=fit_to_block_size(question_input, input_block_size, tokenizer.pad_token_id)
        question_inputs.append(question_input)

        if example.question_varible_output is not None:
            question_varible_output=tokenizer.encode(example.question_varible_output)
        else:
            question_varible_output=tokenizer.build_inputs_with_special_tokens([])
        question_varible_output=fit_to_block_size(question_varible_output, output_block_size, tokenizer.pad_token_id)
        question_varible_outputs.append(question_varible_output)

        if example.condition_output is not None:
            condition_output=tokenizer.encode(example.condition_output)
        else:
            condition_output=tokenizer.build_inputs_with_special_tokens([])
        condition_output=fit_to_block_size(condition_output, output_block_size, tokenizer.pad_token_id)
        condition_outputs.append(condition_output)

    question_inputs = torch.tensor(question_inputs)
    question_varible_outputs = torch.tensor(question_varible_outputs)
    condition_outputs = torch.tensor(condition_outputs)

    question_inputs_mask = build_mask(question_inputs, tokenizer.pad_token_id)
    question_varible_outputs_mask = build_mask(question_varible_outputs, tokenizer.pad_token_id)
    condition_outputs_mask = build_mask(condition_outputs, tokenizer.pad_token_id)

    question_varible_outputs_mask_lm_labels = build_lm_labels(question_varible_outputs, tokenizer.pad_token_id)
    condition_outputs_mask_lm_labels = build_lm_labels(condition_outputs, tokenizer.pad_token_id)

    return (
        question_inputs,
        [question_varible_outputs,condition_outputs],
        question_inputs_mask,
        [question_varible_outputs_mask,condition_outputs_mask],
        [question_varible_outputs_mask_lm_labels,condition_outputs_mask_lm_labels],
    )


# ----------
# Optimizers
# ----------


class BertSumOptimizer(object):
    """ Specific optimizer for BertSum.

    As described in [1], the authors fine-tune BertSum for abstractive
    summarization using two Adam Optimizers with different warm-up steps and
    learning rate. They also use a custom learning rate scheduler.

    [1] Liu, Yang, and Mirella Lapata. "Text summarization with pretrained encoders."
        arXiv preprint arXiv:1908.08345 (2019).
    """

    def __init__(self, model, lr, warmup_steps, beta_1=0.99, beta_2=0.999, eps=1e-8):
        self.encoder = model.encoder
        self.decoders = model.decoders
        self.lr = lr
        self.warmup_steps = warmup_steps
        self.decoders_parameters=[]
        for decoder in model.decoders:
            self.decoders_parameters+=decoder.parameters()

        self.optimizers = {
            "encoder": Adam(
                model.encoder.parameters(),
                lr=lr["encoder"],
                betas=(beta_1, beta_2),
                eps=eps,
            ),
            "decoder": Adam(
                self.decoders_parameters,
                lr=lr["decoder"],
                betas=(beta_1, beta_2),
                eps=eps,
            ),
        }

        self._step = 0

    def _update_rate(self, stack):
        return self.lr[stack] * min(
            self._step ** (-0.5), self._step * self.warmup_steps[stack] ** (-0.5)
        )

    def zero_grad(self):
        self.optimizer_decoder.zero_grad()
        self.optimizer_encoder.zero_grad()

    def step(self):
        self._step += 1
        for stack, optimizer in self.optimizers.items():
            new_rate = self._update_rate(stack)
            for param_group in optimizer.param_groups:
                param_group["lr"] = new_rate
            optimizer.step()


# ------------
# Train
# ------------


def train(args, model, tokenizer):
    """ Fine-tune the pretrained model on the corpus. """
    set_seed(args)

    # Load the data
    args.train_batch_size = args.per_gpu_train_batch_size * max(1, args.n_gpu)
    train_dataset = load_and_cache_examples(args, tokenizer, "train")
    train_sampler = RandomSampler(train_dataset)
    model_collate_fn = functools.partial(collate, tokenizer=tokenizer,
                                         input_block_size=args.input_block_size,output_block_size=args.output_block_size)
    train_dataloader = DataLoader(
        train_dataset,
        sampler=train_sampler,
        batch_size=args.train_batch_size,
        collate_fn=model_collate_fn,
    )

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

    # Prepare the optimizer
    #lr = {"encoder": 0.002, "decoder": 0.2}
    lr = {"encoder": args.encoder_lr, "decoder": args.decoder_lr}
    #warmup_steps = {"encoder": 20000, "decoder": 10000}
    warmup_steps = {"encoder": args.encoder_warmup, "decoder": args.decoder_warmup}
    optimizer = BertSumOptimizer(model, lr, warmup_steps)

    # Train
    logger.info("***** Running training *****")
    logger.info("  Num examples = %d", len(train_dataset))
    logger.info("  Num Epochs = %d", args.num_train_epochs)
    logger.info(
        "  Instantaneous batch size per GPU = %d", args.per_gpu_train_batch_size
    )
    logger.info(
        "  Total train batch size (w. parallel, distributed & accumulation) = %d",
        args.train_batch_size * args.gradient_accumulation_steps
        # * (torch.distributed.get_world_size() if args.local_rank != -1 else 1),
    )
    logger.info("  Gradient Accumulation steps = %d", args.gradient_accumulation_steps)
    logger.info("  Total optimization steps = %d", t_total)

    model.zero_grad()
    train_iterator = trange(args.num_train_epochs, desc="Epoch", disable=False)

    global_step = 0
    tr_loss = 0.0
    for _ in train_iterator:
        epoch_iterator = tqdm(train_dataloader, desc="Iteration", disable=False)
        for step, batch in enumerate(epoch_iterator):
            source, target, encoder_mask, decoder_mask, lm_labels = batch
            #print('source: {}'.format(source))
            #print('target: {}'.format(target))

            feed_source=None
            feed_targets=[None]*len(target)
            feed_encoder_mask=None
            feed_decoder_masks=[None]*len(decoder_mask)
            feed_lm_labels=[None]*len(lm_labels)

            feed_source = source.to(args.device)
            for i in range(len(target)):
                feed_targets[i] = target[i].to(args.device)


            feed_encoder_mask = encoder_mask.to(args.device)
            for i in range(len(decoder_mask)):
                feed_decoder_masks[i] = decoder_mask[i].to(args.device)
            for i in range(len(lm_labels)):
                feed_lm_labels[i] = lm_labels[i].to(args.device)

            model.train()
            #print('debug by zhuoyu: source = {}'.format(source))
            #print('debug by zhuoyu: target = {}'.format(target))
            #print('debug by zhuoyu, device:')
            #print('feed source {}'.format(feed_source.device))
            #print('feed target {}'.format([str(feed_target.device) for feed_target in feed_targets]))
            #print('feed encoder mask {}'.format(feed_encoder_mask.device))
            #print('feed decoder masks {}'.format([str(feed_decoder_mask.device) for feed_decoder_mask in feed_decoder_masks]))
            #print('feed lm labels {}'.format([str(feed_lm_label.device) for feed_lm_label in feed_lm_labels]))
            outputs = model(
                feed_source,
                feed_targets,
                encoder_attention_mask=feed_encoder_mask,
                decoder_attention_mask=feed_decoder_masks,
                decoder_lm_labels=feed_lm_labels,
            )

            loss=0
            for i in range(len(model.decoders)):
                #print('outputs[{}][0] type: {}'.format(i,type(outputs[i][0])))
                loss += outputs[i][0]
            #print(loss)
            if args.gradient_accumulation_steps > 1:
                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()
                model.zero_grad()
                global_step += 1

            if args.max_steps > 0 and global_step > args.max_steps:
                epoch_iterator.close()
                break

        if args.max_steps > 0 and global_step > args.max_steps:
            train_iterator.close()
            break

    return global_step, tr_loss / global_step


# ------------
# Train
# ------------


def evaluate(args, model, tokenizer, prefix=""):
    set_seed(args)

    args.eval_batch_size = args.per_gpu_eval_batch_size * max(1, args.n_gpu)
    eval_dataset = load_and_cache_examples(args, tokenizer, prefix="dev")
    #for example in eval_dataset.examples:
    #    print(example.example_id)
    #    print(example.question_input)
    #    print(example.question_varible_output)
    #    print(example.condition_output)
    #exit(-1)
    eval_sampler = SequentialSampler(eval_dataset)
    model_collate_fn = functools.partial(collate, tokenizer=tokenizer,
                                         input_block_size=args.input_block_size,output_block_size=args.output_block_size)
    eval_dataloader = DataLoader(
        eval_dataset, sampler=eval_sampler, batch_size=args.eval_batch_size,collate_fn=model_collate_fn,
    )

    # multi-gpu evaluate
    #if args.n_gpu > 1:
    #    model = torch.nn.DataParallel(model)

    logger.info("***** Running evaluation {} *****".format(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
    model.eval()
    fout=open(os.path.join(args.output_dir,"dev.res"),'w',encoding='utf-8')
    fdebug=open(os.path.join(args.output_dir,"dev.debug.res"),'w',encoding='utf-8')
    for batch in tqdm(eval_dataloader, desc="Evaluating"):
        source, target, encoder_mask, decoder_mask, lm_labels = batch
        #print('[SOURCE]: {}'.format(source))
        #print('[TARGET]: {}'.format(target))

        #source = source.to(args.device)
        #target = target.to(args.device)

        #encoder_mask = encoder_mask.to(args.device)
        #decoder_mask = decoder_mask.to(args.device)
        #lm_labels = lm_labels.to(args.device)

        feed_source = None
        feed_targets = [None] * len(target)
        feed_encoder_mask = None
        feed_decoder_masks = [None] * len(decoder_mask)
        feed_lm_labels = [None] * len(lm_labels)

        feed_source = source.to(args.device)
        for i in range(len(target)):
            feed_targets[i] = target[i].to(args.device)

        feed_encoder_mask = encoder_mask.to(args.device)
        for i in range(len(decoder_mask)):
            feed_decoder_masks[i] = decoder_mask[i].to(args.device)
        for i in range(len(lm_labels)):
            feed_lm_labels[i] = lm_labels[i].to(args.device)

        with torch.no_grad():

            if args.decoding_type=='decoding':
                tokens_roles=[]
                for i in range(len(feed_targets)):
                    outputs_ids=model.decoding(
                        feed_source,
                        feed_targets[i],
                        encoder_attention_mask=feed_encoder_mask,
                        decoder_attention_mask=feed_decoder_masks[i],
                        decoder_lm_labels=feed_lm_labels[i],
                        decoder=model.decoders[i]
                        #fdebug=fdebug,
                    )
                    print('outputs size: {}'.format(outputs_ids.size()))
                    outputs_ids =outputs_ids.cpu().numpy()


                    batch_tokens=[]
                    for idx in outputs_ids:
                        tokens = []
                        for id in idx:
                            #print('{}\t{}'.format(id,type(id)))
                            tokens.append(tokenizer.ids_to_tokens.get(int(id), tokenizer.unk_token))

                        batch_tokens.append(tokens)

                    tokens_roles.append(batch_tokens)

                def subtoken2token(subtokens):
                    token=""
                    tokens=[]
                    for subtoken in subtokens:
                        if subtoken.startswith("##"):
                            token+=subtoken[2:]
                        else:
                            if token!="":
                                tokens.append(token)
                            token=subtoken
                    if token!="":
                        tokens.append(token)
                    return tokens
                for i in range(len(tokens_roles[0])):
                    fout.write('\t'.join([' '.join(subtoken2token(tokens_roles[0][i]))
                                             ,' '.join(subtoken2token(tokens_roles[1][i]))]) + '\n')

            else:
                print('debug eva input:')
                print('feed_source={}'.format(feed_source))
                print('feed_targets={}'.format(feed_targets))
                print('feed_encoder_mask={}'.format(feed_encoder_mask))
                print('feed_decoder_masks={}'.format(feed_decoder_masks))
                print('feed_lm_labels={}'.format(feed_lm_labels))
                outputs = model(
                    feed_source,
                    feed_targets,
                    encoder_attention_mask=feed_encoder_mask,
                    decoder_attention_mask=feed_decoder_masks,
                    decoder_lm_labels=feed_lm_labels,
                    #fdebug=fdebug,
                )

                ans_seqs=[[],[]]
                for i in range(len(model.decoders)):
                    print(outputs[i][1].size())
                    predicted_scores=outputs[i][1].argmax(-1).cpu().numpy().tolist()
                    for idx in predicted_scores:
                        tokens = []
                        for id in idx:
                            tokens.append(tokenizer.ids_to_tokens.get(id, tokenizer.unk_token))
                        ans_seqs[i].append(tokens)

                for i in range(len(ans_seqs[0])):
                    fout.write('\t'.join([' '.join(ans_seqs[0][i]),' '.join(ans_seqs[1][i])]) + '\n')





                # print('debug by zhuoyu, predicted_scores size={}'.format(predicted_scores.size()))
                #eval_loss += lm_loss.mean().item()




        nb_eval_steps += 1

    eval_loss = eval_loss / nb_eval_steps
    perplexity = torch.exp(torch.tensor(eval_loss))

    result = {"perplexity": perplexity}

    # Save the evaluation's results
    output_eval_file = os.path.join(args.output_dir, "eval_results.txt")
    if not os.path.exists(args.output_dir):
        os.makedirs(args.output_dir)

    with open(output_eval_file, "w") as writer:
        logger.info("***** Eval results {} *****".format(prefix))
        for key in sorted(result.keys()):
            logger.info("  %s = %s", key, str(result[key]))
            writer.write("%s = %s\n" % (key, str(result[key])))

    #with open(os.path.join(args.output_dir,"dev.res"),'w',encoding='utf-8') as fout:
    fout.flush()
    fout.close()
    fdebug.flush()
    fdebug.close()
    return result


def main():
    parser = argparse.ArgumentParser()

    # Required parameters
    parser.add_argument(
        "--data_dir",
        default=None,
        type=str,
        required=True,
        help="The input training data file (a text file).",
    )
    parser.add_argument(
        "--output_dir",
        default=None,
        type=str,
        required=True,
        help="The output directory where the model predictions and checkpoints will be written.",
    )

    # Optional parameters
    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(
        "--do_evaluate",
        type=bool,
        default=False,
        help="Run model evaluation on out-of-sample data.",
    )
    parser.add_argument("--do_train", type=bool, default=False, help="Run training.")
    parser.add_argument(
        "--do_overwrite_output_dir",
        type=bool,
        default=False,
        help="Whether to overwrite the output dir.",
    )
    parser.add_argument(
        "--encoder_model_name_or_path",
        default="bert-base-cased",
        type=str,
        help="The model checkpoint to initialize the encoder's weights with.",
    )
    parser.add_argument(
        "--decoder_model_name_or_path",
        default="/data/zhuoyu/semantic_parsing/models",
        type=str,
        help="The model checkpoint to initialize the decoder's weights with.",
    )
    parser.add_argument(
        "--model_type",
        default="bert",
        type=str,
        help="The decoder architecture to be fine-tuned.",
    )
    parser.add_argument(
        "--max_grad_norm", default=1.0, type=float, help="Max gradient norm."
    )
    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(
        "--to_cpu", default=False, type=bool, help="Whether to force training on CPU."
    )
    parser.add_argument(
        "--num_train_epochs",
        default=10,
        type=int,
        help="Total number of training epochs to perform.",
    )
    parser.add_argument(
        "--per_gpu_eval_batch_size",
        default=4,
        type=int,
        help="Batch size per GPU/CPU for eval.",
    )
    parser.add_argument(
        "--per_gpu_train_batch_size",
        default=4,
        type=int,
        help="Batch size per GPU/CPU for training.",
    )
    parser.add_argument(
        "--input_block_size",
        default=256,
        type=int,
        help="Max seq length for input",
    )
    parser.add_argument(
        "--output_block_size",
        default=64,
        type=int,
        help="Max seq length for output",
    )

    parser.add_argument(
        "--trained_checkpoints",
        default="",
        type=str,
        help="trained_checkpoints",
    )

    parser.add_argument(
        "--decoding_type",
        default="pnt",
        type=str,
        help="",
    )

    parser.add_argument(
        "--encoder_lr",
        default=5e-4,
        type=float,
        help="encoder's learning rate",
    )

    parser.add_argument(
        "--decoder_lr",
        default=5e-4,
        type=float,
        help="encoder's learning rate",
    )

    parser.add_argument(
        "--encoder_warmup",
        default=10,
        type=int,
        help="encoder's learning rate",
    )

    parser.add_argument(
        "--decoder_warmup",
        default=100,
        type=int,
        help="encoder's learning rate",
    )

    parser.add_argument("--seed", default=42, type=int)

    parser.add_argument(
        "--decoder_version",
        default="v1",
        type=str,
        help="",
    )

    args = parser.parse_args()

    if (
                        os.path.exists(args.output_dir)
                    and os.listdir(args.output_dir)
                and args.do_train
            and not args.do_overwrite_output_dir
    ):
        raise ValueError(
            "Output directory ({}) already exists and is not empty. Use --do_overwrite_output_dir to overwrite.".format(
                args.output_dir
            )
        )



    # Set up training device
    if args.to_cpu or not torch.cuda.is_available():
        args.device = torch.device("cpu")
        args.n_gpu = 0
    else:
        args.device = torch.device("cuda")
        args.n_gpu = torch.cuda.device_count()
        print(args.n_gpu)

    # Load pretrained model and tokenizer. The decoder's weights are randomly initialized.
    tokenizer = AutoTokenizer.from_pretrained(args.encoder_model_name_or_path
                                              ,never_split=['[unused0]','[unused1]','[unused2]','[unused3]'])
    #config = BertConfig.from_pretrained(args.model_name_or_path)
    #config.num_hidden_layers=3
    #config.is_decoder=True
    #decoder_model = BertForMaskedLM(config)
    decoder_models=[BertForMaskedLM.from_pretrained(args.decoder_model_name_or_path),
                    BertForMaskedLM.from_pretrained(args.decoder_model_name_or_path)]
    model = Model2Models.from_pretrained(
        args.encoder_model_name_or_path, decoder_model=decoder_models
    )
    #model = Model2Model.from_pretrained(
    #    args.model_name_or_path, decoder_model=None
    #)

    # Setup logging
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        level=logging.INFO,
    )
    logger.warning(
        "Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, 16-bits training: %s",
        0,
        args.device,
        args.n_gpu,
        False,
        False,
    )

    logger.info("Training/evaluation parameters %s", args)

    # Train the model

    if not os.path.exists(args.output_dir):
            os.makedirs(args.output_dir)
    if args.do_train:
        model.to(args.device)
        global_step, tr_loss = train(args, model, tokenizer)
        logger.info(" global_step = %s, average loss = %s", global_step, tr_loss)



        logger.info("Saving model checkpoint to %s", args.output_dir)

        # Save a trained model, configuration and tokenizer using `save_pretrained()`.
        # They can then be reloaded using `from_pretrained()`
        model_to_save = (
            model.module if hasattr(model, "module") else model
        )  # Take care of distributed/parallel training
        model_to_save.save_pretrained(args.output_dir)
        tokenizer.save_pretrained(args.output_dir)
        torch.save(args, os.path.join(args.output_dir, "training_arguments.bin"))

    # Evaluate the model
    results = {}
    if args.do_evaluate:
        checkpoints = [args.trained_checkpoints]
        logger.info("Evaluate the following checkpoints: %s", checkpoints)
        for checkpoint in checkpoints:
            encoder_checkpoint = os.path.join(checkpoint, "encoder")
            decoder_checkpoint_question_varibles = os.path.join(checkpoint, "decoder_0")
            decoder_checkpoint_conditions = os.path.join(checkpoint, "decoder_1")

            decoder_models = [BertForMaskedLM.from_pretrained(decoder_checkpoint_question_varibles),
                              BertForMaskedLM.from_pretrained(decoder_checkpoint_conditions)]
            model = Model2Models.from_pretrained(
                encoder_checkpoint, decoder_model=decoder_models
            )

            model.to(args.device)

            #model = PreTrainedEncoderDecoder.from_pretrained(
            #    encoder_checkpoint, decoder_checkpoint
            #)
            #model = Model2Model.from_pretrained(encoder_checkpoint)
            #model.to(args.device)
            results = "placeholder"

            evaluate(args,model,tokenizer,"test")

    return results


if __name__ == "__main__":
    main()

[((811, 833), 'sys.path.append', 'sys.path.append', (['"""../"""'], {}), "('../')\n", (826, 833), False, 'import sys\n'), ((1699, 1726), 'logging.getLogger', 'logging.getLogger', (['__name__'], {}), '(__name__)\n', (1716, 1726), False, 'import logging\n'), ((1727, 1785), 'logging.basicConfig', 'logging.basicConfig', ([], {'stream': 'sys.stdout', 'level': 'logging.INFO'}), '(stream=sys.stdout, level=logging.INFO)\n', (1746, 1785), False, 'import logging\n'), ((1812, 1834), 'random.seed', 'random.seed', (['args.seed'], {}), '(args.seed)\n', (1823, 1834), False, 'import random\n'), ((1839, 1864), 'numpy.random.seed', 'np.random.seed', (['args.seed'], {}), '(args.seed)\n', (1853, 1864), True, 'import numpy as np\n'), ((1869, 1897), 'torch.manual_seed', 'torch.manual_seed', (['args.seed'], {}), '(args.seed)\n', (1886, 1897), False, 'import torch\n'), ((2023, 2089), 'utils_chemistry.ChemistryDataset', 'ChemistryDataset', (['tokenizer'], {'prefix': 'prefix', 'data_dir': 'args.data_dir'}), '(tokenizer, prefix=prefix, data_dir=args.data_dir)\n', (2039, 2089), False, 'from utils_chemistry import ChemistryDataset\n'), ((3357, 3386), 'torch.tensor', 'torch.tensor', (['question_inputs'], {}), '(question_inputs)\n', (3369, 3386), False, 'import torch\n'), ((3418, 3456), 'torch.tensor', 'torch.tensor', (['question_varible_outputs'], {}), '(question_varible_outputs)\n', (3430, 3456), False, 'import torch\n'), ((3481, 3512), 'torch.tensor', 'torch.tensor', (['condition_outputs'], {}), '(condition_outputs)\n', (3493, 3512), False, 'import torch\n'), ((3541, 3592), 'utils_summarization.build_mask', 'build_mask', (['question_inputs', 'tokenizer.pad_token_id'], {}), '(question_inputs, tokenizer.pad_token_id)\n', (3551, 3592), False, 'from utils_summarization import CNNDailyMailDataset, encode_for_summarization, fit_to_block_size, build_lm_labels, build_mask, compute_token_type_ids\n'), ((3629, 3689), 'utils_summarization.build_mask', 'build_mask', (['question_varible_outputs', 'tokenizer.pad_token_id'], {}), '(question_varible_outputs, tokenizer.pad_token_id)\n', (3639, 3689), False, 'from utils_summarization import CNNDailyMailDataset, encode_for_summarization, fit_to_block_size, build_lm_labels, build_mask, compute_token_type_ids\n'), ((3719, 3772), 'utils_summarization.build_mask', 'build_mask', (['condition_outputs', 'tokenizer.pad_token_id'], {}), '(condition_outputs, tokenizer.pad_token_id)\n', (3729, 3772), False, 'from utils_summarization import CNNDailyMailDataset, encode_for_summarization, fit_to_block_size, build_lm_labels, build_mask, compute_token_type_ids\n'), ((3820, 3885), 'utils_summarization.build_lm_labels', 'build_lm_labels', (['question_varible_outputs', 'tokenizer.pad_token_id'], {}), '(question_varible_outputs, tokenizer.pad_token_id)\n', (3835, 3885), False, 'from utils_summarization import CNNDailyMailDataset, encode_for_summarization, fit_to_block_size, build_lm_labels, build_mask, compute_token_type_ids\n'), ((3925, 3983), 'utils_summarization.build_lm_labels', 'build_lm_labels', (['condition_outputs', 'tokenizer.pad_token_id'], {}), '(condition_outputs, tokenizer.pad_token_id)\n', (3940, 3983), False, 'from utils_summarization import CNNDailyMailDataset, encode_for_summarization, fit_to_block_size, build_lm_labels, build_mask, compute_token_type_ids\n'), ((6429, 6457), 'torch.utils.data.RandomSampler', 'RandomSampler', (['train_dataset'], {}), '(train_dataset)\n', (6442, 6457), False, 'from torch.utils.data import DataLoader, RandomSampler, SequentialSampler\n'), ((6481, 6615), 'functools.partial', 'functools.partial', (['collate'], {'tokenizer': 'tokenizer', 'input_block_size': 'args.input_block_size', 'output_block_size': 'args.output_block_size'}), '(collate, tokenizer=tokenizer, input_block_size=args.\n    input_block_size, output_block_size=args.output_block_size)\n', (6498, 6615), False, 'import functools\n'), ((6674, 6790), 'torch.utils.data.DataLoader', 'DataLoader', (['train_dataset'], {'sampler': 'train_sampler', 'batch_size': 'args.train_batch_size', 'collate_fn': 'model_collate_fn'}), '(train_dataset, sampler=train_sampler, batch_size=args.\n    train_batch_size, collate_fn=model_collate_fn)\n', (6684, 6790), False, 'from torch.utils.data import DataLoader, RandomSampler, SequentialSampler\n'), ((8265, 8323), 'tqdm.trange', 'trange', (['args.num_train_epochs'], {'desc': '"""Epoch"""', 'disable': '(False)'}), "(args.num_train_epochs, desc='Epoch', disable=False)\n", (8271, 8323), False, 'from tqdm import tqdm, trange\n'), ((11748, 11779), 'torch.utils.data.SequentialSampler', 'SequentialSampler', (['eval_dataset'], {}), '(eval_dataset)\n', (11765, 11779), False, 'from torch.utils.data import DataLoader, RandomSampler, SequentialSampler\n'), ((11803, 11937), 'functools.partial', 'functools.partial', (['collate'], {'tokenizer': 'tokenizer', 'input_block_size': 'args.input_block_size', 'output_block_size': 'args.output_block_size'}), '(collate, tokenizer=tokenizer, input_block_size=args.\n    input_block_size, output_block_size=args.output_block_size)\n', (11820, 11937), False, 'import functools\n'), ((11995, 12108), 'torch.utils.data.DataLoader', 'DataLoader', (['eval_dataset'], {'sampler': 'eval_sampler', 'batch_size': 'args.eval_batch_size', 'collate_fn': 'model_collate_fn'}), '(eval_dataset, sampler=eval_sampler, batch_size=args.\n    eval_batch_size, collate_fn=model_collate_fn)\n', (12005, 12108), False, 'from torch.utils.data import DataLoader, RandomSampler, SequentialSampler\n'), ((12636, 12676), 'tqdm.tqdm', 'tqdm', (['eval_dataloader'], {'desc': '"""Evaluating"""'}), "(eval_dataloader, desc='Evaluating')\n", (12640, 12676), False, 'from tqdm import tqdm, trange\n'), ((17392, 17441), 'os.path.join', 'os.path.join', (['args.output_dir', '"""eval_results.txt"""'], {}), "(args.output_dir, 'eval_results.txt')\n", (17404, 17441), False, 'import os\n'), ((18008, 18033), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ([], {}), '()\n', (18031, 18033), False, 'import argparse\n'), ((22658, 22791), 'transformers.AutoTokenizer.from_pretrained', 'AutoTokenizer.from_pretrained', (['args.encoder_model_name_or_path'], {'never_split': "['[unused0]', '[unused1]', '[unused2]', '[unused3]']"}), "(args.encoder_model_name_or_path, never_split=\n    ['[unused0]', '[unused1]', '[unused2]', '[unused3]'])\n", (22687, 22791), False, 'from transformers import AutoTokenizer, BertForMaskedLM, BertConfig, PreTrainedEncoderDecoder, Model2Models\n'), ((23185, 23281), 'transformers.Model2Models.from_pretrained', 'Model2Models.from_pretrained', (['args.encoder_model_name_or_path'], {'decoder_model': 'decoder_models'}), '(args.encoder_model_name_or_path, decoder_model\n    =decoder_models)\n', (23213, 23281), False, 'from transformers import AutoTokenizer, BertForMaskedLM, BertConfig, PreTrainedEncoderDecoder, Model2Models\n'), ((23418, 23561), 'logging.basicConfig', 'logging.basicConfig', ([], {'format': '"""%(asctime)s - %(levelname)s - %(name)s -   %(message)s"""', 'datefmt': '"""%m/%d/%Y %H:%M:%S"""', 'level': 'logging.INFO'}), "(format=\n    '%(asctime)s - %(levelname)s - %(name)s -   %(message)s', datefmt=\n    '%m/%d/%Y %H:%M:%S', level=logging.INFO)\n", (23437, 23561), False, 'import logging\n'), ((2421, 2496), 'utils_summarization.fit_to_block_size', 'fit_to_block_size', (['question_input', 'input_block_size', 'tokenizer.pad_token_id'], {}), '(question_input, input_block_size, tokenizer.pad_token_id)\n', (2438, 2496), False, 'from utils_summarization import CNNDailyMailDataset, encode_for_summarization, fit_to_block_size, build_lm_labels, build_mask, compute_token_type_ids\n'), ((2816, 2906), 'utils_summarization.fit_to_block_size', 'fit_to_block_size', (['question_varible_output', 'output_block_size', 'tokenizer.pad_token_id'], {}), '(question_varible_output, output_block_size, tokenizer.\n    pad_token_id)\n', (2833, 2906), False, 'from utils_summarization import CNNDailyMailDataset, encode_for_summarization, fit_to_block_size, build_lm_labels, build_mask, compute_token_type_ids\n'), ((3204, 3282), 'utils_summarization.fit_to_block_size', 'fit_to_block_size', (['condition_output', 'output_block_size', 'tokenizer.pad_token_id'], {}), '(condition_output, output_block_size, tokenizer.pad_token_id)\n', (3221, 3282), False, 'from utils_summarization import CNNDailyMailDataset, encode_for_summarization, fit_to_block_size, build_lm_labels, build_mask, compute_token_type_ids\n'), ((8417, 8472), 'tqdm.tqdm', 'tqdm', (['train_dataloader'], {'desc': '"""Iteration"""', 'disable': '(False)'}), "(train_dataloader, desc='Iteration', disable=False)\n", (8421, 8472), False, 'from tqdm import tqdm, trange\n'), ((12473, 12513), 'os.path.join', 'os.path.join', (['args.output_dir', '"""dev.res"""'], {}), "(args.output_dir, 'dev.res')\n", (12485, 12513), False, 'import os\n'), ((12551, 12597), 'os.path.join', 'os.path.join', (['args.output_dir', '"""dev.debug.res"""'], {}), "(args.output_dir, 'dev.debug.res')\n", (12563, 12597), False, 'import os\n'), ((17266, 17289), 'torch.tensor', 'torch.tensor', (['eval_loss'], {}), '(eval_loss)\n', (17278, 17289), False, 'import torch\n'), ((17453, 17484), 'os.path.exists', 'os.path.exists', (['args.output_dir'], {}), '(args.output_dir)\n', (17467, 17484), False, 'import os\n'), ((17494, 17522), 'os.makedirs', 'os.makedirs', (['args.output_dir'], {}), '(args.output_dir)\n', (17505, 17522), False, 'import os\n'), ((21897, 21928), 'os.path.exists', 'os.path.exists', (['args.output_dir'], {}), '(args.output_dir)\n', (21911, 21928), False, 'import os\n'), ((21953, 21980), 'os.listdir', 'os.listdir', (['args.output_dir'], {}), '(args.output_dir)\n', (21963, 21980), False, 'import os\n'), ((22381, 22400), 'torch.device', 'torch.device', (['"""cpu"""'], {}), "('cpu')\n", (22393, 22400), False, 'import torch\n'), ((22456, 22476), 'torch.device', 'torch.device', (['"""cuda"""'], {}), "('cuda')\n", (22468, 22476), False, 'import torch\n'), ((22498, 22523), 'torch.cuda.device_count', 'torch.cuda.device_count', ([], {}), '()\n', (22521, 22523), False, 'import torch\n'), ((23021, 23085), 'transformers.BertForMaskedLM.from_pretrained', 'BertForMaskedLM.from_pretrained', (['args.decoder_model_name_or_path'], {}), '(args.decoder_model_name_or_path)\n', (23052, 23085), False, 'from transformers import AutoTokenizer, BertForMaskedLM, BertConfig, PreTrainedEncoderDecoder, Model2Models\n'), ((23107, 23171), 'transformers.BertForMaskedLM.from_pretrained', 'BertForMaskedLM.from_pretrained', (['args.decoder_model_name_or_path'], {}), '(args.decoder_model_name_or_path)\n', (23138, 23171), False, 'from transformers import AutoTokenizer, BertForMaskedLM, BertConfig, PreTrainedEncoderDecoder, Model2Models\n'), ((23885, 23916), 'os.path.exists', 'os.path.exists', (['args.output_dir'], {}), '(args.output_dir)\n', (23899, 23916), False, 'import os\n'), ((23930, 23958), 'os.makedirs', 'os.makedirs', (['args.output_dir'], {}), '(args.output_dir)\n', (23941, 23958), False, 'import os\n'), ((5333, 5418), 'torch.optim.Adam', 'Adam', (['self.decoders_parameters'], {'lr': "lr['decoder']", 'betas': '(beta_1, beta_2)', 'eps': 'eps'}), "(self.decoders_parameters, lr=lr['decoder'], betas=(beta_1, beta_2),\n    eps=eps)\n", (5337, 5418), False, 'from torch.optim import Adam\n'), ((13710, 13725), 'torch.no_grad', 'torch.no_grad', ([], {}), '()\n', (13723, 13725), False, 'import torch\n'), ((22332, 22357), 'torch.cuda.is_available', 'torch.cuda.is_available', ([], {}), '()\n', (22355, 22357), False, 'import torch\n'), ((24654, 24709), 'os.path.join', 'os.path.join', (['args.output_dir', '"""training_arguments.bin"""'], {}), "(args.output_dir, 'training_arguments.bin')\n", (24666, 24709), False, 'import os\n'), ((24975, 25010), 'os.path.join', 'os.path.join', (['checkpoint', '"""encoder"""'], {}), "(checkpoint, 'encoder')\n", (24987, 25010), False, 'import os\n'), ((25062, 25099), 'os.path.join', 'os.path.join', (['checkpoint', '"""decoder_0"""'], {}), "(checkpoint, 'decoder_0')\n", (25074, 25099), False, 'import os\n'), ((25144, 25181), 'os.path.join', 'os.path.join', (['checkpoint', '"""decoder_1"""'], {}), "(checkpoint, 'decoder_1')\n", (25156, 25181), False, 'import os\n'), ((25398, 25476), 'transformers.Model2Models.from_pretrained', 'Model2Models.from_pretrained', (['encoder_checkpoint'], {'decoder_model': 'decoder_models'}), '(encoder_checkpoint, decoder_model=decoder_models)\n', (25426, 25476), False, 'from transformers import AutoTokenizer, BertForMaskedLM, BertConfig, PreTrainedEncoderDecoder, Model2Models\n'), ((25213, 25282), 'transformers.BertForMaskedLM.from_pretrained', 'BertForMaskedLM.from_pretrained', (['decoder_checkpoint_question_varibles'], {}), '(decoder_checkpoint_question_varibles)\n', (25244, 25282), False, 'from transformers import AutoTokenizer, BertForMaskedLM, BertConfig, PreTrainedEncoderDecoder, Model2Models\n'), ((25314, 25376), 'transformers.BertForMaskedLM.from_pretrained', 'BertForMaskedLM.from_pretrained', (['decoder_checkpoint_conditions'], {}), '(decoder_checkpoint_conditions)\n', (25345, 25376), False, 'from transformers import AutoTokenizer, BertForMaskedLM, BertConfig, PreTrainedEncoderDecoder, Model2Models\n')]