chendl/compositional_test

1#!/usr/bin/env python2# coding=utf-83# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.4# Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.5#6# Licensed under the Apache License, Version 2.0 (the "License");7# you may not use this file except in compliance with the License.8# You may obtain a copy of the License at9#10#     http://www.apache.org/licenses/LICENSE-2.011#12# Unless required by applicable law or agreed to in writing, software13# distributed under the License is distributed on an "AS IS" BASIS,14# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.15# See the License for the specific language governing permissions and16# limitations under the License.17"""18Fine-tuning the library models for language modeling on a text file (GPT, GPT-2, CTRL, BERT, RoBERTa, XLNet).19GPT, GPT-2 and CTRL are fine-tuned using a causal language modeling (CLM) loss. BERT and RoBERTa are fine-tuned20using a masked language modeling (MLM) loss. XLNet is fine-tuned using a permutation language modeling (PLM) loss.21"""22 23 24import logging25import math26import os27from dataclasses import dataclass, field28from glob import glob29from typing import Optional30 31from torch.utils.data import ConcatDataset32 33import transformers34from transformers import (35    CONFIG_MAPPING,36    MODEL_WITH_LM_HEAD_MAPPING,37    AutoConfig,38    AutoModelWithLMHead,39    AutoTokenizer,40    DataCollatorForLanguageModeling,41    DataCollatorForPermutationLanguageModeling,42    DataCollatorForWholeWordMask,43    HfArgumentParser,44    LineByLineTextDataset,45    LineByLineWithRefDataset,46    PreTrainedTokenizer,47    TextDataset,48    Trainer,49    TrainingArguments,50    set_seed,51)52from transformers.trainer_utils import is_main_process53 54 55logger = logging.getLogger(__name__)56 57 58MODEL_CONFIG_CLASSES = list(MODEL_WITH_LM_HEAD_MAPPING.keys())59MODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)60 61 62@dataclass63class ModelArguments:64    """65    Arguments pertaining to which model/config/tokenizer we are going to fine-tune, or train from scratch.66    """67 68    model_name_or_path: Optional[str] = field(69        default=None,70        metadata={71            "help": (72                "The model checkpoint for weights initialization. Leave None if you want to train a model from"73                " scratch."74            )75        },76    )77    model_type: Optional[str] = field(78        default=None,79        metadata={"help": "If training from scratch, pass a model type from the list: " + ", ".join(MODEL_TYPES)},80    )81    config_name: Optional[str] = field(82        default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}83    )84    tokenizer_name: Optional[str] = field(85        default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}86    )87    cache_dir: Optional[str] = field(88        default=None,89        metadata={"help": "Where do you want to store the pretrained models downloaded from huggingface.co"},90    )91 92 93@dataclass94class DataTrainingArguments:95    """96    Arguments pertaining to what data we are going to input our model for training and eval.97    """98 99    train_data_file: Optional[str] = field(100        default=None, metadata={"help": "The input training data file (a text file)."}101    )102    train_data_files: Optional[str] = field(103        default=None,104        metadata={105            "help": (106                "The input training data files (multiple files in glob format). "107                "Very often splitting large files to smaller files can prevent tokenizer going out of memory"108            )109        },110    )111    eval_data_file: Optional[str] = field(112        default=None,113        metadata={"help": "An optional input evaluation data file to evaluate the perplexity on (a text file)."},114    )115    train_ref_file: Optional[str] = field(116        default=None,117        metadata={"help": "An optional input train ref data file for whole word mask in Chinese."},118    )119    eval_ref_file: Optional[str] = field(120        default=None,121        metadata={"help": "An optional input eval ref data file for whole word mask in Chinese."},122    )123    line_by_line: bool = field(124        default=False,125        metadata={"help": "Whether distinct lines of text in the dataset are to be handled as distinct sequences."},126    )127 128    mlm: bool = field(129        default=False, metadata={"help": "Train with masked-language modeling loss instead of language modeling."}130    )131    whole_word_mask: bool = field(default=False, metadata={"help": "Whether ot not to use whole word mask."})132    mlm_probability: float = field(133        default=0.15, metadata={"help": "Ratio of tokens to mask for masked language modeling loss"}134    )135    plm_probability: float = field(136        default=1 / 6,137        metadata={138            "help": (139                "Ratio of length of a span of masked tokens to surrounding context length for permutation language"140                " modeling."141            )142        },143    )144    max_span_length: int = field(145        default=5, metadata={"help": "Maximum length of a span of masked tokens for permutation language modeling."}146    )147 148    block_size: int = field(149        default=-1,150        metadata={151            "help": (152                "Optional input sequence length after tokenization."153                "The training dataset will be truncated in block of this size for training."154                "Default to the model max input length for single sentence inputs (take into account special tokens)."155            )156        },157    )158    overwrite_cache: bool = field(159        default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}160    )161 162 163def get_dataset(164    args: DataTrainingArguments,165    tokenizer: PreTrainedTokenizer,166    evaluate: bool = False,167    cache_dir: Optional[str] = None,168):169    def _dataset(file_path, ref_path=None):170        if args.line_by_line:171            if ref_path is not None:172                if not args.whole_word_mask or not args.mlm:173                    raise ValueError("You need to set world whole masking and mlm to True for Chinese Whole Word Mask")174                return LineByLineWithRefDataset(175                    tokenizer=tokenizer,176                    file_path=file_path,177                    block_size=args.block_size,178                    ref_path=ref_path,179                )180 181            return LineByLineTextDataset(tokenizer=tokenizer, file_path=file_path, block_size=args.block_size)182        else:183            return TextDataset(184                tokenizer=tokenizer,185                file_path=file_path,186                block_size=args.block_size,187                overwrite_cache=args.overwrite_cache,188                cache_dir=cache_dir,189            )190 191    if evaluate:192        return _dataset(args.eval_data_file, args.eval_ref_file)193    elif args.train_data_files:194        return ConcatDataset([_dataset(f) for f in glob(args.train_data_files)])195    else:196        return _dataset(args.train_data_file, args.train_ref_file)197 198 199def main():200    # See all possible arguments in src/transformers/training_args.py201    # or by passing the --help flag to this script.202    # We now keep distinct sets of args, for a cleaner separation of concerns.203 204    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))205    model_args, data_args, training_args = parser.parse_args_into_dataclasses()206 207    if data_args.eval_data_file is None and training_args.do_eval:208        raise ValueError(209            "Cannot do evaluation without an evaluation data file. Either supply a file to --eval_data_file "210            "or remove the --do_eval argument."211        )212    if (213        os.path.exists(training_args.output_dir)214        and os.listdir(training_args.output_dir)215        and training_args.do_train216        and not training_args.overwrite_output_dir217    ):218        raise ValueError(219            f"Output directory ({training_args.output_dir}) already exists and is not empty. Use"220            " --overwrite_output_dir to overcome."221        )222 223    # Setup logging224    logging.basicConfig(225        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",226        datefmt="%m/%d/%Y %H:%M:%S",227        level=logging.INFO if training_args.local_rank in [-1, 0] else logging.WARN,228    )229    logger.warning(230        "Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, 16-bits training: %s",231        training_args.local_rank,232        training_args.device,233        training_args.n_gpu,234        bool(training_args.local_rank != -1),235        training_args.fp16,236    )237    # Set the verbosity to info of the Transformers logger (on main process only):238    if is_main_process(training_args.local_rank):239        transformers.utils.logging.set_verbosity_info()240        transformers.utils.logging.enable_default_handler()241        transformers.utils.logging.enable_explicit_format()242    logger.info("Training/evaluation parameters %s", training_args)243 244    # Set seed245    set_seed(training_args.seed)246 247    # Load pretrained model and tokenizer248    #249    # Distributed training:250    # The .from_pretrained methods guarantee that only one local process can concurrently251    # download model & vocab.252 253    if model_args.config_name:254        config = AutoConfig.from_pretrained(model_args.config_name, cache_dir=model_args.cache_dir)255    elif model_args.model_name_or_path:256        config = AutoConfig.from_pretrained(model_args.model_name_or_path, cache_dir=model_args.cache_dir)257    else:258        config = CONFIG_MAPPING[model_args.model_type]()259        logger.warning("You are instantiating a new config instance from scratch.")260 261    if model_args.tokenizer_name:262        tokenizer = AutoTokenizer.from_pretrained(model_args.tokenizer_name, cache_dir=model_args.cache_dir)263    elif model_args.model_name_or_path:264        tokenizer = AutoTokenizer.from_pretrained(model_args.model_name_or_path, cache_dir=model_args.cache_dir)265    else:266        raise ValueError(267            "You are instantiating a new tokenizer from scratch. This is not supported, but you can do it from another"268            " script, save it,and load it from here, using --tokenizer_name"269        )270 271    if model_args.model_name_or_path:272        model = AutoModelWithLMHead.from_pretrained(273            model_args.model_name_or_path,274            from_tf=bool(".ckpt" in model_args.model_name_or_path),275            config=config,276            cache_dir=model_args.cache_dir,277        )278    else:279        logger.info("Training new model from scratch")280        model = AutoModelWithLMHead.from_config(config)281 282    model.resize_token_embeddings(len(tokenizer))283 284    if config.model_type in ["bert", "roberta", "distilbert", "camembert"] and not data_args.mlm:285        raise ValueError(286            "BERT and RoBERTa-like models do not have LM heads but masked LM heads. They must be run using the"287            "--mlm flag (masked language modeling)."288        )289 290    if data_args.block_size <= 0:291        data_args.block_size = tokenizer.max_len292        # Our input block size will be the max possible for the model293    else:294        data_args.block_size = min(data_args.block_size, tokenizer.max_len)295 296    # Get datasets297 298    train_dataset = (299        get_dataset(data_args, tokenizer=tokenizer, cache_dir=model_args.cache_dir) if training_args.do_train else None300    )301    eval_dataset = (302        get_dataset(data_args, tokenizer=tokenizer, evaluate=True, cache_dir=model_args.cache_dir)303        if training_args.do_eval304        else None305    )306    if config.model_type == "xlnet":307        data_collator = DataCollatorForPermutationLanguageModeling(308            tokenizer=tokenizer,309            plm_probability=data_args.plm_probability,310            max_span_length=data_args.max_span_length,311        )312    else:313        if data_args.mlm and data_args.whole_word_mask:314            data_collator = DataCollatorForWholeWordMask(315                tokenizer=tokenizer, mlm_probability=data_args.mlm_probability316            )317        else:318            data_collator = DataCollatorForLanguageModeling(319                tokenizer=tokenizer, mlm=data_args.mlm, mlm_probability=data_args.mlm_probability320            )321 322    # Initialize our Trainer323    trainer = Trainer(324        model=model,325        args=training_args,326        data_collator=data_collator,327        train_dataset=train_dataset,328        eval_dataset=eval_dataset,329        prediction_loss_only=True,330    )331 332    # Training333    if training_args.do_train:334        model_path = (335            model_args.model_name_or_path336            if model_args.model_name_or_path is not None and os.path.isdir(model_args.model_name_or_path)337            else None338        )339        trainer.train(model_path=model_path)340        trainer.save_model()341        # For convenience, we also re-save the tokenizer to the same directory,342        # so that you can share your model easily on huggingface.co/models =)343        if trainer.is_world_master():344            tokenizer.save_pretrained(training_args.output_dir)345 346    # Evaluation347    results = {}348    if training_args.do_eval:349        logger.info("*** Evaluate ***")350 351        eval_output = trainer.evaluate()352 353        perplexity = math.exp(eval_output["eval_loss"])354        result = {"perplexity": perplexity}355 356        output_eval_file = os.path.join(training_args.output_dir, "eval_results_lm.txt")357        if trainer.is_world_master():358            with open(output_eval_file, "w") as writer:359                logger.info("***** Eval results *****")360                for key in sorted(result.keys()):361                    logger.info("  %s = %s", key, str(result[key]))362                    writer.write("%s = %s\n" % (key, str(result[key])))363 364        results.update(result)365 366    return results367 368 369def _mp_fn(index):370    # For xla_spawn (TPUs)371    main()372 373 374if __name__ == "__main__":375    main()376