Aluode/PerceptionLabPortable

1# Copyright 2020 The HuggingFace Team. All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7#     http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14 15import os16import time17import warnings18from dataclasses import dataclass, field19from enum import Enum20from typing import Optional, Union21 22import torch23from filelock import FileLock24from torch.utils.data import Dataset25 26from ...tokenization_utils_base import PreTrainedTokenizerBase27from ...utils import check_torch_load_is_safe, logging28from ..processors.glue import glue_convert_examples_to_features, glue_output_modes, glue_processors29from ..processors.utils import InputFeatures30 31 32logger = logging.get_logger(__name__)33 34 35@dataclass36class GlueDataTrainingArguments:37    """38    Arguments pertaining to what data we are going to input our model for training and eval.39 40    Using `HfArgumentParser` we can turn this class into argparse arguments to be able to specify them on the command41    line.42    """43 44    task_name: str = field(metadata={"help": "The name of the task to train on: " + ", ".join(glue_processors.keys())})45    data_dir: str = field(46        metadata={"help": "The input data dir. Should contain the .tsv files (or other data files) for the task."}47    )48    max_seq_length: int = field(49        default=128,50        metadata={51            "help": (52                "The maximum total input sequence length after tokenization. Sequences longer "53                "than this will be truncated, sequences shorter will be padded."54            )55        },56    )57    overwrite_cache: bool = field(58        default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}59    )60 61    def __post_init__(self):62        self.task_name = self.task_name.lower()63 64 65class Split(Enum):66    train = "train"67    dev = "dev"68    test = "test"69 70 71class GlueDataset(Dataset):72    """73    This will be superseded by a framework-agnostic approach soon.74    """75 76    args: GlueDataTrainingArguments77    output_mode: str78    features: list[InputFeatures]79 80    def __init__(81        self,82        args: GlueDataTrainingArguments,83        tokenizer: PreTrainedTokenizerBase,84        limit_length: Optional[int] = None,85        mode: Union[str, Split] = Split.train,86        cache_dir: Optional[str] = None,87    ):88        warnings.warn(89            "This dataset will be removed from the library soon, preprocessing should be handled with the 🤗 Datasets "90            "library. You can have a look at this example script for pointers: "91            "https://github.com/huggingface/transformers/blob/main/examples/pytorch/text-classification/run_glue.py",92            FutureWarning,93        )94        self.args = args95        self.processor = glue_processors[args.task_name]()96        self.output_mode = glue_output_modes[args.task_name]97        if isinstance(mode, str):98            try:99                mode = Split[mode]100            except KeyError:101                raise KeyError("mode is not a valid split name")102        # Load data features from cache or dataset file103        cached_features_file = os.path.join(104            cache_dir if cache_dir is not None else args.data_dir,105            f"cached_{mode.value}_{tokenizer.__class__.__name__}_{args.max_seq_length}_{args.task_name}",106        )107        label_list = self.processor.get_labels()108        if args.task_name in ["mnli", "mnli-mm"] and tokenizer.__class__.__name__ in (109            "RobertaTokenizer",110            "RobertaTokenizerFast",111            "XLMRobertaTokenizer",112            "BartTokenizer",113            "BartTokenizerFast",114        ):115            # HACK(label indices are swapped in RoBERTa pretrained model)116            label_list[1], label_list[2] = label_list[2], label_list[1]117        self.label_list = label_list118 119        # Make sure only the first process in distributed training processes the dataset,120        # and the others will use the cache.121        lock_path = cached_features_file + ".lock"122        with FileLock(lock_path):123            if os.path.exists(cached_features_file) and not args.overwrite_cache:124                start = time.time()125                check_torch_load_is_safe()126                self.features = torch.load(cached_features_file, weights_only=True)127                logger.info(128                    f"Loading features from cached file {cached_features_file} [took %.3f s]", time.time() - start129                )130            else:131                logger.info(f"Creating features from dataset file at {args.data_dir}")132 133                if mode == Split.dev:134                    examples = self.processor.get_dev_examples(args.data_dir)135                elif mode == Split.test:136                    examples = self.processor.get_test_examples(args.data_dir)137                else:138                    examples = self.processor.get_train_examples(args.data_dir)139                if limit_length is not None:140                    examples = examples[:limit_length]141                self.features = glue_convert_examples_to_features(142                    examples,143                    tokenizer,144                    max_length=args.max_seq_length,145                    label_list=label_list,146                    output_mode=self.output_mode,147                )148                start = time.time()149                torch.save(self.features, cached_features_file)150                # ^ This seems to take a lot of time so I want to investigate why and how we can improve.151                logger.info(152                    f"Saving features into cached file {cached_features_file} [took {time.time() - start:.3f} s]"153                )154 155    def __len__(self):156        return len(self.features)157 158    def __getitem__(self, i) -> InputFeatures:159        return self.features[i]160 161    def get_labels(self):162        return self.label_list163