chendl/compositional_test

1#!/usr/bin/env python2# coding=utf-83# Copyright 2021 The HuggingFace Inc. team. All rights reserved.4#5# Licensed under the Apache License, Version 2.0 (the "License");6# you may not use this file except in compliance with the License.7# You may obtain a copy of the License at8#9#     http://www.apache.org/licenses/LICENSE-2.010#11# Unless required by applicable law or agreed to in writing, software12# distributed under the License is distributed on an "AS IS" BASIS,13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.14# See the License for the specific language governing permissions and15# limitations under the License.16 17""" Fine-tuning a 🤗 Transformers CTC model for automatic speech recognition"""18 19import functools20import json21import logging22import os23import re24import sys25import warnings26from dataclasses import dataclass, field27from typing import Dict, List, Optional, Union28 29import datasets30import evaluate31import numpy as np32import torch33from datasets import DatasetDict, load_dataset34 35import transformers36from transformers import (37    AutoConfig,38    AutoFeatureExtractor,39    AutoModelForCTC,40    AutoProcessor,41    AutoTokenizer,42    HfArgumentParser,43    Trainer,44    TrainingArguments,45    Wav2Vec2Processor,46    set_seed,47)48from transformers.trainer_utils import get_last_checkpoint, is_main_process49from transformers.utils import check_min_version, send_example_telemetry50from transformers.utils.versions import require_version51 52 53# Will error if the minimal version of Transformers is not installed. Remove at your own risks.54check_min_version("4.28.0")55 56require_version("datasets>=1.18.0", "To fix: pip install -r examples/pytorch/speech-recognition/requirements.txt")57 58 59logger = logging.getLogger(__name__)60 61 62def list_field(default=None, metadata=None):63    return field(default_factory=lambda: default, metadata=metadata)64 65 66@dataclass67class ModelArguments:68    """69    Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.70    """71 72    model_name_or_path: str = field(73        metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}74    )75    tokenizer_name_or_path: Optional[str] = field(76        default=None,77        metadata={"help": "Path to pretrained tokenizer or tokenizer identifier from huggingface.co/models"},78    )79    cache_dir: Optional[str] = field(80        default=None,81        metadata={"help": "Where do you want to store the pretrained models downloaded from huggingface.co"},82    )83    freeze_feature_encoder: bool = field(84        default=True, metadata={"help": "Whether to freeze the feature encoder layers of the model."}85    )86    attention_dropout: float = field(87        default=0.0, metadata={"help": "The dropout ratio for the attention probabilities."}88    )89    activation_dropout: float = field(90        default=0.0, metadata={"help": "The dropout ratio for activations inside the fully connected layer."}91    )92    feat_proj_dropout: float = field(default=0.0, metadata={"help": "The dropout ratio for the projected features."})93    hidden_dropout: float = field(94        default=0.0,95        metadata={96            "help": "The dropout probability for all fully connected layers in the embeddings, encoder, and pooler."97        },98    )99    final_dropout: float = field(100        default=0.0,101        metadata={"help": "The dropout probability for the final projection layer."},102    )103    mask_time_prob: float = field(104        default=0.05,105        metadata={106            "help": (107                "Probability of each feature vector along the time axis to be chosen as the start of the vector"108                "span to be masked. Approximately ``mask_time_prob * sequence_length // mask_time_length`` feature"109                "vectors will be masked along the time axis."110            )111        },112    )113    mask_time_length: int = field(114        default=10,115        metadata={"help": "Length of vector span to mask along the time axis."},116    )117    mask_feature_prob: float = field(118        default=0.0,119        metadata={120            "help": (121                "Probability of each feature vector along the feature axis to be chosen as the start of the vectorspan"122                " to be masked. Approximately ``mask_feature_prob * sequence_length // mask_feature_length`` feature"123                " bins will be masked along the time axis."124            )125        },126    )127    mask_feature_length: int = field(128        default=10,129        metadata={"help": "Length of vector span to mask along the feature axis."},130    )131    layerdrop: float = field(default=0.0, metadata={"help": "The LayerDrop probability."})132    ctc_loss_reduction: Optional[str] = field(133        default="mean", metadata={"help": "The way the ctc loss should be reduced. Should be one of 'mean' or 'sum'."}134    )135 136 137@dataclass138class DataTrainingArguments:139    """140    Arguments pertaining to what data we are going to input our model for training and eval.141 142    Using `HfArgumentParser` we can turn this class143    into argparse arguments to be able to specify them on144    the command line.145    """146 147    dataset_name: str = field(148        metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}149    )150    dataset_config_name: str = field(151        default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}152    )153    train_split_name: str = field(154        default="train+validation",155        metadata={156            "help": (157                "The name of the training data set split to use (via the datasets library). Defaults to "158                "'train+validation'"159            )160        },161    )162    eval_split_name: str = field(163        default="test",164        metadata={165            "help": "The name of the evaluation data set split to use (via the datasets library). Defaults to 'test'"166        },167    )168    audio_column_name: str = field(169        default="audio",170        metadata={"help": "The name of the dataset column containing the audio data. Defaults to 'audio'"},171    )172    text_column_name: str = field(173        default="text",174        metadata={"help": "The name of the dataset column containing the text data. Defaults to 'text'"},175    )176    overwrite_cache: bool = field(177        default=False, metadata={"help": "Overwrite the cached preprocessed datasets or not."}178    )179    preprocessing_num_workers: Optional[int] = field(180        default=None,181        metadata={"help": "The number of processes to use for the preprocessing."},182    )183    max_train_samples: Optional[int] = field(184        default=None,185        metadata={186            "help": (187                "For debugging purposes or quicker training, truncate the number of training examples to this "188                "value if set."189            )190        },191    )192    max_eval_samples: Optional[int] = field(193        default=None,194        metadata={195            "help": (196                "For debugging purposes or quicker training, truncate the number of validation examples to this "197                "value if set."198            )199        },200    )201    chars_to_ignore: Optional[List[str]] = list_field(202        default=None,203        metadata={"help": "A list of characters to remove from the transcripts."},204    )205    eval_metrics: List[str] = list_field(206        default=["wer"],207        metadata={"help": "A list of metrics the model should be evaluated on. E.g. `'wer cer'`"},208    )209    max_duration_in_seconds: float = field(210        default=20.0,211        metadata={212            "help": (213                "Filter audio files that are longer than `max_duration_in_seconds` seconds to"214                " 'max_duration_in_seconds`"215            )216        },217    )218    min_duration_in_seconds: float = field(219        default=0.0, metadata={"help": "Filter audio files that are shorter than `min_duration_in_seconds` seconds"}220    )221    preprocessing_only: bool = field(222        default=False,223        metadata={224            "help": (225                "Whether to only do data preprocessing and skip training. This is especially useful when data"226                " preprocessing errors out in distributed training due to timeout. In this case, one should run the"227                " preprocessing in a non-distributed setup with `preprocessing_only=True` so that the cached datasets"228                " can consequently be loaded in distributed training"229            )230        },231    )232    use_auth_token: bool = field(233        default=False,234        metadata={235            "help": (236                "If :obj:`True`, will use the token generated when running"237                ":obj:`huggingface-cli login` as HTTP bearer authorization for remote files."238            )239        },240    )241    unk_token: str = field(242        default="[UNK]",243        metadata={"help": "The unk token for the tokenizer"},244    )245    pad_token: str = field(246        default="[PAD]",247        metadata={"help": "The padding token for the tokenizer"},248    )249    word_delimiter_token: str = field(250        default="|",251        metadata={"help": "The word delimiter token for the tokenizer"},252    )253    phoneme_language: Optional[str] = field(254        default=None,255        metadata={256            "help": (257                "The target language that should be used be"258                " passed to the tokenizer for tokenization. Note that"259                " this is only relevant if the model classifies the"260                " input audio to a sequence of phoneme sequences."261            )262        },263    )264 265 266@dataclass267class DataCollatorCTCWithPadding:268    """269    Data collator that will dynamically pad the inputs received.270    Args:271        processor (:class:`~transformers.AutoProcessor`)272            The processor used for proccessing the data.273        padding (:obj:`bool`, :obj:`str` or :class:`~transformers.tokenization_utils_base.PaddingStrategy`, `optional`, defaults to :obj:`True`):274            Select a strategy to pad the returned sequences (according to the model's padding side and padding index)275            among:276            * :obj:`True` or :obj:`'longest'`: Pad to the longest sequence in the batch (or no padding if only a single277              sequence if provided).278            * :obj:`'max_length'`: Pad to a maximum length specified with the argument :obj:`max_length` or to the279              maximum acceptable input length for the model if that argument is not provided.280            * :obj:`False` or :obj:`'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of281              different lengths).282        max_length (:obj:`int`, `optional`):283            Maximum length of the ``input_values`` of the returned list and optionally padding length (see above).284        max_length_labels (:obj:`int`, `optional`):285            Maximum length of the ``labels`` returned list and optionally padding length (see above).286        pad_to_multiple_of (:obj:`int`, `optional`):287            If set will pad the sequence to a multiple of the provided value.288            This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability >=289            7.5 (Volta).290    """291 292    processor: AutoProcessor293    padding: Union[bool, str] = "longest"294    pad_to_multiple_of: Optional[int] = None295    pad_to_multiple_of_labels: Optional[int] = None296 297    def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]:298        # split inputs and labels since they have to be of different lenghts and need299        # different padding methods300        input_features = [{"input_values": feature["input_values"]} for feature in features]301        label_features = [{"input_ids": feature["labels"]} for feature in features]302 303        batch = self.processor.pad(304            input_features,305            padding=self.padding,306            pad_to_multiple_of=self.pad_to_multiple_of,307            return_tensors="pt",308        )309 310        labels_batch = self.processor.pad(311            labels=label_features,312            padding=self.padding,313            pad_to_multiple_of=self.pad_to_multiple_of_labels,314            return_tensors="pt",315        )316 317        # replace padding with -100 to ignore loss correctly318        labels = labels_batch["input_ids"].masked_fill(labels_batch.attention_mask.ne(1), -100)319 320        batch["labels"] = labels321        if "attention_mask" in batch:322            batch["attention_mask"] = batch["attention_mask"].to(torch.long)323 324        return batch325 326 327def create_vocabulary_from_data(328    datasets: DatasetDict,329    word_delimiter_token: Optional[str] = None,330    unk_token: Optional[str] = None,331    pad_token: Optional[str] = None,332):333    # Given training and test labels create vocabulary334    def extract_all_chars(batch):335        all_text = " ".join(batch["target_text"])336        vocab = list(set(all_text))337        return {"vocab": [vocab], "all_text": [all_text]}338 339    vocabs = datasets.map(340        extract_all_chars,341        batched=True,342        batch_size=-1,343        keep_in_memory=True,344        remove_columns=datasets["train"].column_names,345    )346 347    # take union of all unique characters in each dataset348    vocab_set = functools.reduce(349        lambda vocab_1, vocab_2: set(vocab_1["vocab"][0]) | set(vocab_2["vocab"][0]), vocabs.values()350    )351 352    vocab_dict = {v: k for k, v in enumerate(sorted(vocab_set))}353 354    # replace white space with delimiter token355    if word_delimiter_token is not None:356        vocab_dict[word_delimiter_token] = vocab_dict[" "]357        del vocab_dict[" "]358 359    # add unk and pad token360    if unk_token is not None:361        vocab_dict[unk_token] = len(vocab_dict)362 363    if pad_token is not None:364        vocab_dict[pad_token] = len(vocab_dict)365 366    return vocab_dict367 368 369def main():370    # See all possible arguments in src/transformers/training_args.py371    # or by passing the --help flag to this script.372    # We now keep distinct sets of args, for a cleaner separation of concerns.373 374    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))375    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):376        # If we pass only one argument to the script and it's the path to a json file,377        # let's parse it to get our arguments.378        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))379    else:380        model_args, data_args, training_args = parser.parse_args_into_dataclasses()381 382    # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The383    # information sent is the one passed as arguments along with your Python/PyTorch versions.384    send_example_telemetry("run_speech_recognition_ctc", model_args, data_args)385 386    # Detecting last checkpoint.387    last_checkpoint = None388    if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir:389        last_checkpoint = get_last_checkpoint(training_args.output_dir)390        if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:391            raise ValueError(392                f"Output directory ({training_args.output_dir}) already exists and is not empty. "393                "Use --overwrite_output_dir to overcome."394            )395        elif last_checkpoint is not None:396            logger.info(397                f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "398                "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."399            )400 401    # Setup logging402    logging.basicConfig(403        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",404        datefmt="%m/%d/%Y %H:%M:%S",405        handlers=[logging.StreamHandler(sys.stdout)],406    )407    logger.setLevel(logging.INFO if is_main_process(training_args.local_rank) else logging.WARN)408 409    # Log on each process the small summary:410    logger.warning(411        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"412        f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"413    )414    # Set the verbosity to info of the Transformers logger (on main process only):415    if is_main_process(training_args.local_rank):416        transformers.utils.logging.set_verbosity_info()417    logger.info("Training/evaluation parameters %s", training_args)418 419    # Set seed before initializing model.420    set_seed(training_args.seed)421 422    # 1. First, let's load the dataset423    raw_datasets = DatasetDict()424 425    if training_args.do_train:426        raw_datasets["train"] = load_dataset(427            data_args.dataset_name,428            data_args.dataset_config_name,429            split=data_args.train_split_name,430            use_auth_token=data_args.use_auth_token,431        )432 433        if data_args.audio_column_name not in raw_datasets["train"].column_names:434            raise ValueError(435                f"--audio_column_name '{data_args.audio_column_name}' not found in dataset '{data_args.dataset_name}'."436                " Make sure to set `--audio_column_name` to the correct audio column - one of"437                f" {', '.join(raw_datasets['train'].column_names)}."438            )439 440        if data_args.text_column_name not in raw_datasets["train"].column_names:441            raise ValueError(442                f"--text_column_name {data_args.text_column_name} not found in dataset '{data_args.dataset_name}'. "443                "Make sure to set `--text_column_name` to the correct text column - one of "444                f"{', '.join(raw_datasets['train'].column_names)}."445            )446 447        if data_args.max_train_samples is not None:448            raw_datasets["train"] = raw_datasets["train"].select(range(data_args.max_train_samples))449 450    if training_args.do_eval:451        raw_datasets["eval"] = load_dataset(452            data_args.dataset_name,453            data_args.dataset_config_name,454            split=data_args.eval_split_name,455            use_auth_token=data_args.use_auth_token,456        )457 458        if data_args.max_eval_samples is not None:459            raw_datasets["eval"] = raw_datasets["eval"].select(range(data_args.max_eval_samples))460 461    # 2. We remove some special characters from the datasets462    # that make training complicated and do not help in transcribing the speech463    # E.g. characters, such as `,` and `.` do not really have an acoustic characteristic464    # that could be easily picked up by the model465    chars_to_ignore_regex = (466        f'[{"".join(data_args.chars_to_ignore)}]' if data_args.chars_to_ignore is not None else None467    )468    text_column_name = data_args.text_column_name469 470    def remove_special_characters(batch):471        if chars_to_ignore_regex is not None:472            batch["target_text"] = re.sub(chars_to_ignore_regex, "", batch[text_column_name]).lower() + " "473        else:474            batch["target_text"] = batch[text_column_name].lower() + " "475        return batch476 477    with training_args.main_process_first(desc="dataset map special characters removal"):478        raw_datasets = raw_datasets.map(479            remove_special_characters,480            remove_columns=[text_column_name],481            desc="remove special characters from datasets",482        )483 484    # save special tokens for tokenizer485    word_delimiter_token = data_args.word_delimiter_token486    unk_token = data_args.unk_token487    pad_token = data_args.pad_token488 489    # 3. Next, let's load the config as we might need it to create490    # the tokenizer491    # load config492    config = AutoConfig.from_pretrained(493        model_args.model_name_or_path, cache_dir=model_args.cache_dir, use_auth_token=data_args.use_auth_token494    )495 496    # 4. Next, if no tokenizer file is defined,497    # we create the vocabulary of the model by extracting all unique characters from498    # the training and evaluation datasets499    # We need to make sure that only first rank saves vocabulary500    # make sure all processes wait until vocab is created501    tokenizer_name_or_path = model_args.tokenizer_name_or_path502    tokenizer_kwargs = {}503    if tokenizer_name_or_path is None:504        # save vocab in training output dir505        tokenizer_name_or_path = training_args.output_dir506 507        vocab_file = os.path.join(tokenizer_name_or_path, "vocab.json")508 509        with training_args.main_process_first():510            if training_args.overwrite_output_dir and os.path.isfile(vocab_file):511                try:512                    os.remove(vocab_file)513                except OSError:514                    # in shared file-systems it might be the case that515                    # two processes try to delete the vocab file at the some time516                    pass517 518        with training_args.main_process_first(desc="dataset map vocabulary creation"):519            if not os.path.isfile(vocab_file):520                os.makedirs(tokenizer_name_or_path, exist_ok=True)521                vocab_dict = create_vocabulary_from_data(522                    raw_datasets,523                    word_delimiter_token=word_delimiter_token,524                    unk_token=unk_token,525                    pad_token=pad_token,526                )527 528                # save vocab dict to be loaded into tokenizer529                with open(vocab_file, "w") as file:530                    json.dump(vocab_dict, file)531 532        # if tokenizer has just been created533        # it is defined by `tokenizer_class` if present in config else by `model_type`534        tokenizer_kwargs = {535            "config": config if config.tokenizer_class is not None else None,536            "tokenizer_type": config.model_type if config.tokenizer_class is None else None,537            "unk_token": unk_token,538            "pad_token": pad_token,539            "word_delimiter_token": word_delimiter_token,540        }541 542    # 5. Now we can instantiate the feature extractor, tokenizer and model543    # Note for distributed training, the .from_pretrained methods guarantee that only544    # one local process can concurrently download model & vocab.545 546    # load feature_extractor and tokenizer547    tokenizer = AutoTokenizer.from_pretrained(548        tokenizer_name_or_path,549        use_auth_token=data_args.use_auth_token,550        **tokenizer_kwargs,551    )552    feature_extractor = AutoFeatureExtractor.from_pretrained(553        model_args.model_name_or_path, cache_dir=model_args.cache_dir, use_auth_token=data_args.use_auth_token554    )555 556    # adapt config557    config.update(558        {559            "feat_proj_dropout": model_args.feat_proj_dropout,560            "attention_dropout": model_args.attention_dropout,561            "hidden_dropout": model_args.hidden_dropout,562            "final_dropout": model_args.final_dropout,563            "mask_time_prob": model_args.mask_time_prob,564            "mask_time_length": model_args.mask_time_length,565            "mask_feature_prob": model_args.mask_feature_prob,566            "mask_feature_length": model_args.mask_feature_length,567            "gradient_checkpointing": training_args.gradient_checkpointing,568            "layerdrop": model_args.layerdrop,569            "ctc_loss_reduction": model_args.ctc_loss_reduction,570            "pad_token_id": tokenizer.pad_token_id,571            "vocab_size": len(tokenizer),572            "activation_dropout": model_args.activation_dropout,573        }574    )575 576    # create model577    model = AutoModelForCTC.from_pretrained(578        model_args.model_name_or_path,579        cache_dir=model_args.cache_dir,580        config=config,581        use_auth_token=data_args.use_auth_token,582    )583 584    # freeze encoder585    if model_args.freeze_feature_encoder:586        model.freeze_feature_encoder()587 588    # 6. Now we preprocess the datasets including loading the audio, resampling and normalization589    # Thankfully, `datasets` takes care of automatically loading and resampling the audio,590    # so that we just need to set the correct target sampling rate and normalize the input591    # via the `feature_extractor`592 593    # make sure that dataset decodes audio with correct sampling rate594    dataset_sampling_rate = next(iter(raw_datasets.values())).features[data_args.audio_column_name].sampling_rate595    if dataset_sampling_rate != feature_extractor.sampling_rate:596        raw_datasets = raw_datasets.cast_column(597            data_args.audio_column_name, datasets.features.Audio(sampling_rate=feature_extractor.sampling_rate)598        )599 600    # derive max & min input length for sample rate & max duration601    max_input_length = data_args.max_duration_in_seconds * feature_extractor.sampling_rate602    min_input_length = data_args.min_duration_in_seconds * feature_extractor.sampling_rate603    audio_column_name = data_args.audio_column_name604    num_workers = data_args.preprocessing_num_workers605 606    # `phoneme_language` is only relevant if the model is fine-tuned on phoneme classification607    phoneme_language = data_args.phoneme_language608 609    # Preprocessing the datasets.610    # We need to read the audio files as arrays and tokenize the targets.611    def prepare_dataset(batch):612        # load audio613        sample = batch[audio_column_name]614 615        inputs = feature_extractor(sample["array"], sampling_rate=sample["sampling_rate"])616        batch["input_values"] = inputs.input_values[0]617        batch["input_length"] = len(batch["input_values"])618 619        # encode targets620        additional_kwargs = {}621        if phoneme_language is not None:622            additional_kwargs["phonemizer_lang"] = phoneme_language623 624        batch["labels"] = tokenizer(batch["target_text"], **additional_kwargs).input_ids625        return batch626 627    with training_args.main_process_first(desc="dataset map preprocessing"):628        vectorized_datasets = raw_datasets.map(629            prepare_dataset,630            remove_columns=next(iter(raw_datasets.values())).column_names,631            num_proc=num_workers,632            desc="preprocess datasets",633        )634 635        def is_audio_in_length_range(length):636            return length > min_input_length and length < max_input_length637 638        # filter data that is shorter than min_input_length639        vectorized_datasets = vectorized_datasets.filter(640            is_audio_in_length_range,641            num_proc=num_workers,642            input_columns=["input_length"],643        )644 645    # 7. Next, we can prepare the training.646    # Let's use word error rate (WER) as our evaluation metric,647    # instantiate a data collator and the trainer648 649    # Define evaluation metrics during training, *i.e.* word error rate, character error rate650    eval_metrics = {metric: evaluate.load(metric) for metric in data_args.eval_metrics}651 652    # for large datasets it is advised to run the preprocessing on a653    # single machine first with ``args.preprocessing_only`` since there will mostly likely654    # be a timeout when running the script in distributed mode.655    # In a second step ``args.preprocessing_only`` can then be set to `False` to load the656    # cached dataset657    if data_args.preprocessing_only:658        logger.info(f"Data preprocessing finished. Files cached at {vectorized_datasets.cache_files}")659        return660 661    def compute_metrics(pred):662        pred_logits = pred.predictions663        pred_ids = np.argmax(pred_logits, axis=-1)664 665        pred.label_ids[pred.label_ids == -100] = tokenizer.pad_token_id666 667        pred_str = tokenizer.batch_decode(pred_ids)668        # we do not want to group tokens when computing the metrics669        label_str = tokenizer.batch_decode(pred.label_ids, group_tokens=False)670 671        metrics = {k: v.compute(predictions=pred_str, references=label_str) for k, v in eval_metrics.items()}672 673        return metrics674 675    # Now save everything to be able to create a single processor later676    # make sure all processes wait until data is saved677    with training_args.main_process_first():678        # only the main process saves them679        if is_main_process(training_args.local_rank):680            # save feature extractor, tokenizer and config681            feature_extractor.save_pretrained(training_args.output_dir)682            tokenizer.save_pretrained(training_args.output_dir)683            config.save_pretrained(training_args.output_dir)684 685    try:686        processor = AutoProcessor.from_pretrained(training_args.output_dir)687    except (OSError, KeyError):688        warnings.warn(689            "Loading a processor from a feature extractor config that does not"690            " include a `processor_class` attribute is deprecated and will be removed in v5. Please add the following "691            " attribute to your `preprocessor_config.json` file to suppress this warning: "692            " `'processor_class': 'Wav2Vec2Processor'`",693            FutureWarning,694        )695        processor = Wav2Vec2Processor.from_pretrained(training_args.output_dir)696 697    # Instantiate custom data collator698    data_collator = DataCollatorCTCWithPadding(processor=processor)699 700    # Initialize Trainer701    trainer = Trainer(702        model=model,703        data_collator=data_collator,704        args=training_args,705        compute_metrics=compute_metrics,706        train_dataset=vectorized_datasets["train"] if training_args.do_train else None,707        eval_dataset=vectorized_datasets["eval"] if training_args.do_eval else None,708        tokenizer=feature_extractor,709    )710 711    # 8. Finally, we can start training712 713    # Training714    if training_args.do_train:715        # use last checkpoint if exist716        if last_checkpoint is not None:717            checkpoint = last_checkpoint718        elif os.path.isdir(model_args.model_name_or_path):719            checkpoint = model_args.model_name_or_path720        else:721            checkpoint = None722 723        train_result = trainer.train(resume_from_checkpoint=checkpoint)724        trainer.save_model()725 726        metrics = train_result.metrics727        max_train_samples = (728            data_args.max_train_samples729            if data_args.max_train_samples is not None730            else len(vectorized_datasets["train"])731        )732        metrics["train_samples"] = min(max_train_samples, len(vectorized_datasets["train"]))733 734        trainer.log_metrics("train", metrics)735        trainer.save_metrics("train", metrics)736        trainer.save_state()737 738    # Evaluation739    results = {}740    if training_args.do_eval:741        logger.info("*** Evaluate ***")742        metrics = trainer.evaluate()743        max_eval_samples = (744            data_args.max_eval_samples if data_args.max_eval_samples is not None else len(vectorized_datasets["eval"])745        )746        metrics["eval_samples"] = min(max_eval_samples, len(vectorized_datasets["eval"]))747 748        trainer.log_metrics("eval", metrics)749        trainer.save_metrics("eval", metrics)750 751    # Write model card and (optionally) push to hub752    config_name = data_args.dataset_config_name if data_args.dataset_config_name is not None else "na"753    kwargs = {754        "finetuned_from": model_args.model_name_or_path,755        "tasks": "automatic-speech-recognition",756        "tags": ["automatic-speech-recognition", data_args.dataset_name],757        "dataset_args": (758            f"Config: {config_name}, Training split: {data_args.train_split_name}, Eval split:"759            f" {data_args.eval_split_name}"760        ),761        "dataset": f"{data_args.dataset_name.upper()} - {config_name.upper()}",762    }763    if "common_voice" in data_args.dataset_name:764        kwargs["language"] = config_name765 766    if training_args.push_to_hub:767        trainer.push_to_hub(**kwargs)768    else:769        trainer.create_model_card(**kwargs)770 771    return results772 773 774if __name__ == "__main__":775    main()776