CoolFace
Modelpublic

Plim/test_lm

sourceHugging Faceapache-2.0updated 5y agoView on Hugging Face
0likes21downloads
run_speech_recognition_ctc.py735 linesDownload Raw Back to root
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 16""" Fine-tuning a ๐Ÿค— Transformers CTC model for automatic speech recognition"""17 18import functools19import json20import logging21import os22import re23import sys24import warnings25from dataclasses import dataclass, field26from typing import Dict, List, Optional, Union27 28import datasets29import numpy as np30import torch31from datasets import DatasetDict, load_dataset, load_metric32 33import transformers34from transformers import (35    AutoConfig,36    AutoFeatureExtractor,37    AutoModelForCTC,38    AutoProcessor,39    AutoTokenizer,40    HfArgumentParser,41    Trainer,42    TrainingArguments,43    Wav2Vec2Processor,44    set_seed,45)46from transformers.trainer_utils import get_last_checkpoint, is_main_process47from transformers.utils import check_min_version48from transformers.utils.versions import require_version49 50 51# Will error if the minimal version of Transformers is not installed. Remove at your own risks.52check_min_version("4.16.0.dev0")53 54require_version("datasets>=1.13.3", "To fix: pip install -r examples/pytorch/text-classification/requirements.txt")55 56 57logger = logging.getLogger(__name__)58 59 60def list_field(default=None, metadata=None):61    return field(default_factory=lambda: default, metadata=metadata)62 63 64@dataclass65class ModelArguments:66    """67    Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.68    """69 70    model_name_or_path: str = field(71        metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}72    )73    tokenizer_name_or_path: Optional[str] = field(74        default=None,75        metadata={"help": "Path to pretrained tokenizer or tokenizer identifier from huggingface.co/models"},76    )77    cache_dir: Optional[str] = field(78        default=None,79        metadata={"help": "Where do you want to store the pretrained models downloaded from huggingface.co"},80    )81    freeze_feature_encoder: bool = field(82        default=True, metadata={"help": "Whether to freeze the feature encoder layers of the model."}83    )84    attention_dropout: float = field(85        default=0.0, metadata={"help": "The dropout ratio for the attention probabilities."}86    )87    activation_dropout: float = field(88        default=0.0, metadata={"help": "The dropout ratio for activations inside the fully connected layer."}89    )90    feat_proj_dropout: float = field(default=0.0, metadata={"help": "The dropout ratio for the projected features."})91    hidden_dropout: float = field(92        default=0.0,93        metadata={94            "help": "The dropout probability for all fully connected layers in the embeddings, encoder, and pooler."95        },96    )97    final_dropout: float = field(98        default=0.0,99        metadata={"help": "The dropout probability for the final projection layer."},100    )101    mask_time_prob: float = field(102        default=0.05,103        metadata={104            "help": "Probability of each feature vector along the time axis to be chosen as the start of the vector"105            "span to be masked. Approximately ``mask_time_prob * sequence_length // mask_time_length`` feature"106            "vectors will be masked along the time axis."107        },108    )109    mask_time_length: int = field(110        default=10,111        metadata={"help": "Length of vector span to mask along the time axis."},112    )113    mask_feature_prob: float = field(114        default=0.0,115        metadata={116            "help": "Probability of each feature vector along the feature axis to be chosen as the start of the vector"117            "span to be masked. Approximately ``mask_feature_prob * sequence_length // mask_feature_length`` feature bins will be masked along the time axis."118        },119    )120    mask_feature_length: int = field(121        default=10,122        metadata={"help": "Length of vector span to mask along the feature axis."},123    )124    layerdrop: float = field(default=0.0, metadata={"help": "The LayerDrop probability."})125    ctc_loss_reduction: Optional[str] = field(126        default="mean", metadata={"help": "The way the ctc loss should be reduced. Should be one of 'mean' or 'sum'."}127    )128 129 130@dataclass131class DataTrainingArguments:132    """133    Arguments pertaining to what data we are going to input our model for training and eval.134 135    Using `HfArgumentParser` we can turn this class136    into argparse arguments to be able to specify them on137    the command line.138    """139 140    dataset_name: str = field(141        metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}142    )143    dataset_config_name: str = field(144        default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}145    )146    train_split_name: str = field(147        default="train+validation",148        metadata={149            "help": "The name of the training data set split to use (via the datasets library). Defaults to 'train'"150        },151    )152    eval_split_name: str = field(153        default="test",154        metadata={155            "help": "The name of the training data set split to use (via the datasets library). Defaults to 'test'"156        },157    )158    audio_column_name: str = field(159        default="audio",160        metadata={"help": "The name of the dataset column containing the audio data. Defaults to 'audio'"},161    )162    text_column_name: str = field(163        default="text",164        metadata={"help": "The name of the dataset column containing the text data. Defaults to 'text'"},165    )166    overwrite_cache: bool = field(167        default=False, metadata={"help": "Overwrite the cached preprocessed datasets or not."}168    )169    preprocessing_num_workers: Optional[int] = field(170        default=None,171        metadata={"help": "The number of processes to use for the preprocessing."},172    )173    max_train_samples: Optional[int] = field(174        default=None,175        metadata={176            "help": "For debugging purposes or quicker training, truncate the number of training examples to this "177            "value if set."178        },179    )180    max_eval_samples: Optional[int] = field(181        default=None,182        metadata={183            "help": "For debugging purposes or quicker training, truncate the number of validation examples to this "184            "value if set."185        },186    )187    chars_to_ignore: Optional[List[str]] = list_field(188        default=None,189        metadata={"help": "A list of characters to remove from the transcripts."},190    )191    eval_metrics: List[str] = list_field(192        default=["wer"],193        metadata={"help": "A list of metrics the model should be evaluated on. E.g. `'wer cer'`"},194    )195    max_duration_in_seconds: float = field(196        default=20.0,197        metadata={198            "help": "Filter audio files that are longer than `max_duration_in_seconds` seconds to 'max_duration_in_seconds`"199        },200    )201    min_duration_in_seconds: float = field(202        default=0.0, metadata={"help": "Filter audio files that are shorter than `min_duration_in_seconds` seconds"}203    )204    preprocessing_only: bool = field(205        default=False,206        metadata={207            "help": "Whether to only do data preprocessing and skip training. "208            "This is especially useful when data preprocessing errors out in distributed training due to timeout. "209            "In this case, one should run the preprocessing in a non-distributed setup with `preprocessing_only=True` "210            "so that the cached datasets can consequently be loaded in distributed training"211        },212    )213    use_auth_token: bool = field(214        default=False,215        metadata={216            "help": "If :obj:`True`, will use the token generated when running"217            ":obj:`transformers-cli login` as HTTP bearer authorization for remote files."218        },219    )220    unk_token: str = field(221        default="[UNK]",222        metadata={"help": "The unk token for the tokenizer"},223    )224    pad_token: str = field(225        default="[PAD]",226        metadata={"help": "The padding token for the tokenizer"},227    )228    word_delimiter_token: str = field(229        default="|",230        metadata={"help": "The word delimiter token for the tokenizer"},231    )232    phoneme_language: Optional[str] = field(233        default=None,234        metadata={235            "help": "The target language that should be used be"236            " passed to the tokenizer for tokenization. Note that"237            " this is only relevant if the model classifies the"238            " input audio to a sequence of phoneme sequences."239        },240    )241 242 243@dataclass244class DataCollatorCTCWithPadding:245    """246    Data collator that will dynamically pad the inputs received.247    Args:248        processor (:class:`~transformers.AutoProcessor`)249            The processor used for proccessing the data.250        padding (:obj:`bool`, :obj:`str` or :class:`~transformers.tokenization_utils_base.PaddingStrategy`, `optional`, defaults to :obj:`True`):251            Select a strategy to pad the returned sequences (according to the model's padding side and padding index)252            among:253            * :obj:`True` or :obj:`'longest'`: Pad to the longest sequence in the batch (or no padding if only a single254              sequence if provided).255            * :obj:`'max_length'`: Pad to a maximum length specified with the argument :obj:`max_length` or to the256              maximum acceptable input length for the model if that argument is not provided.257            * :obj:`False` or :obj:`'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of258              different lengths).259        max_length (:obj:`int`, `optional`):260            Maximum length of the ``input_values`` of the returned list and optionally padding length (see above).261        max_length_labels (:obj:`int`, `optional`):262            Maximum length of the ``labels`` returned list and optionally padding length (see above).263        pad_to_multiple_of (:obj:`int`, `optional`):264            If set will pad the sequence to a multiple of the provided value.265            This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability >=266            7.5 (Volta).267    """268 269    processor: AutoProcessor270    padding: Union[bool, str] = "longest"271    pad_to_multiple_of: Optional[int] = None272    pad_to_multiple_of_labels: Optional[int] = None273 274    def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]:275        # split inputs and labels since they have to be of different lenghts and need276        # different padding methods277        input_features = [{"input_values": feature["input_values"]} for feature in features]278        label_features = [{"input_ids": feature["labels"]} for feature in features]279 280        batch = self.processor.pad(281            input_features,282            padding=self.padding,283            pad_to_multiple_of=self.pad_to_multiple_of,284            return_tensors="pt",285        )286 287        with self.processor.as_target_processor():288            labels_batch = self.processor.pad(289                label_features,290                padding=self.padding,291                pad_to_multiple_of=self.pad_to_multiple_of_labels,292                return_tensors="pt",293            )294 295        # replace padding with -100 to ignore loss correctly296        labels = labels_batch["input_ids"].masked_fill(labels_batch.attention_mask.ne(1), -100)297 298        batch["labels"] = labels299 300        return batch301 302 303def create_vocabulary_from_data(304    datasets: DatasetDict,305    word_delimiter_token: Optional[str] = None,306    unk_token: Optional[str] = None,307    pad_token: Optional[str] = None,308):309    # Given training and test labels create vocabulary310    def extract_all_chars(batch):311        all_text = " ".join(batch["target_text"])312        vocab = list(set(all_text))313        return {"vocab": [vocab], "all_text": [all_text]}314 315    vocabs = datasets.map(316        extract_all_chars,317        batched=True,318        batch_size=-1,319        keep_in_memory=True,320        remove_columns=datasets["train"].column_names,321    )322 323    # take union of all unique characters in each dataset324    vocab_set = functools.reduce(325        lambda vocab_1, vocab_2: set(vocab_1["vocab"][0]) | set(vocab_2["vocab"][0]), vocabs.values()326    )327 328    vocab_dict = {v: k for k, v in enumerate(sorted(list(vocab_set)))}329 330    # replace white space with delimiter token331    if word_delimiter_token is not None:332        vocab_dict[word_delimiter_token] = vocab_dict[" "]333        del vocab_dict[" "]334 335    # add unk and pad token336    if unk_token is not None:337        vocab_dict[unk_token] = len(vocab_dict)338 339    if pad_token is not None:340        vocab_dict[pad_token] = len(vocab_dict)341 342    return vocab_dict343 344 345def main():346    # See all possible arguments in src/transformers/training_args.py347    # or by passing the --help flag to this script.348    # We now keep distinct sets of args, for a cleaner separation of concerns.349 350    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))351    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):352        # If we pass only one argument to the script and it's the path to a json file,353        # let's parse it to get our arguments.354        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))355    else:356        model_args, data_args, training_args = parser.parse_args_into_dataclasses()357 358    # Detecting last checkpoint.359    last_checkpoint = None360    if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir:361        last_checkpoint = get_last_checkpoint(training_args.output_dir)362        if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:363            raise ValueError(364                f"Output directory ({training_args.output_dir}) already exists and is not empty. "365                "Use --overwrite_output_dir to overcome."366            )367        elif last_checkpoint is not None:368            logger.info(369                f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "370                "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."371            )372 373    # Setup logging374    logging.basicConfig(375        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",376        datefmt="%m/%d/%Y %H:%M:%S",377        handlers=[logging.StreamHandler(sys.stdout)],378    )379    logger.setLevel(logging.INFO if is_main_process(training_args.local_rank) else logging.WARN)380 381    # Log on each process the small summary:382    logger.warning(383        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"384        f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"385    )386    # Set the verbosity to info of the Transformers logger (on main process only):387    if is_main_process(training_args.local_rank):388        transformers.utils.logging.set_verbosity_info()389    logger.info("Training/evaluation parameters %s", training_args)390 391    # Set seed before initializing model.392    set_seed(training_args.seed)393 394    # 1. First, let's load the dataset395    raw_datasets = DatasetDict()396 397    if training_args.do_train:398        raw_datasets["train"] = load_dataset(399            data_args.dataset_name,400            data_args.dataset_config_name,401            split=data_args.train_split_name,402            use_auth_token=data_args.use_auth_token,403        )404 405        if data_args.audio_column_name not in raw_datasets["train"].column_names:406            raise ValueError(407                f"--audio_column_name '{data_args.audio_column_name}' not found in dataset '{data_args.dataset_name}'. "408                "Make sure to set `--audio_column_name` to the correct audio column - one of "409                f"{', '.join(raw_datasets['train'].column_names)}."410            )411 412        if data_args.text_column_name not in raw_datasets["train"].column_names:413            raise ValueError(414                f"--text_column_name {data_args.text_column_name} not found in dataset '{data_args.dataset_name}'. "415                "Make sure to set `--text_column_name` to the correct text column - one of "416                f"{', '.join(raw_datasets['train'].column_names)}."417            )418 419        if data_args.max_train_samples is not None:420            raw_datasets["train"] = raw_datasets["train"].select(range(data_args.max_train_samples))421 422    if training_args.do_eval:423        raw_datasets["eval"] = load_dataset(424            data_args.dataset_name,425            data_args.dataset_config_name,426            split=data_args.eval_split_name,427            use_auth_token=data_args.use_auth_token,428        )429 430        if data_args.max_eval_samples is not None:431            raw_datasets["eval"] = raw_datasets["eval"].select(range(data_args.max_eval_samples))432 433    # 2. We remove some special characters from the datasets434    # that make training complicated and do not help in transcribing the speech435    # E.g. characters, such as `,` and `.` do not really have an acoustic characteristic436    # that could be easily picked up by the model437    chars_to_ignore_regex = '[^a-zร รขรครงรฉรจรชรซรฎรฏรดรถรนรปรผรฟ\'โ€™ ]'438    text_column_name = data_args.text_column_name439 440    def remove_and_replace_special_characters(batch):441        batch["target_text"] = re.sub(chars_to_ignore_regex, "", batch[text_column_name].lower()).replace('โ€™', "'") + " "442        return batch443 444    with training_args.main_process_first(desc="dataset map special characters removal"):445        raw_datasets = raw_datasets.map(446            remove_and_replace_special_characters,447            remove_columns=[text_column_name],448            desc="remove special characters from datasets",449        )450 451    # save special tokens for tokenizer452    word_delimiter_token = data_args.word_delimiter_token453    unk_token = data_args.unk_token454    pad_token = data_args.pad_token455 456    # 3. Next, let's load the config as we might need it to create457    # the tokenizer458    # load config459    config = AutoConfig.from_pretrained(460        model_args.model_name_or_path, cache_dir=model_args.cache_dir, use_auth_token=data_args.use_auth_token461    )462 463    # 4. Next, if no tokenizer file is defined,464    # we create the vocabulary of the model by extracting all unique characters from465    # the training and evaluation datasets466    # We need to make sure that only first rank saves vocabulary467    # make sure all processes wait until vocab is created468    tokenizer_name_or_path = model_args.tokenizer_name_or_path469    tokenizer_kwargs = {}470    if tokenizer_name_or_path is None:471        # save vocab in training output dir472        tokenizer_name_or_path = training_args.output_dir473 474        vocab_file = os.path.join(tokenizer_name_or_path, "vocab.json")475 476        with training_args.main_process_first():477            if training_args.overwrite_output_dir and os.path.isfile(vocab_file):478                os.remove(vocab_file)479 480        with training_args.main_process_first(desc="dataset map vocabulary creation"):481            if not os.path.isfile(vocab_file):482                os.makedirs(tokenizer_name_or_path, exist_ok=True)483                vocab_dict = create_vocabulary_from_data(484                    raw_datasets,485                    word_delimiter_token=word_delimiter_token,486                    unk_token=unk_token,487                    pad_token=pad_token,488                )489 490                # save vocab dict to be loaded into tokenizer491                with open(vocab_file, "w") as file:492                    json.dump(vocab_dict, file)493 494        # if tokenizer has just been created495        # it is defined by `tokenizer_class` if present in config else by `model_type`496        tokenizer_kwargs = {497            "config": config if config.tokenizer_class is not None else None,498            "tokenizer_type": config.model_type if config.tokenizer_class is None else None,499            "unk_token": unk_token,500            "pad_token": pad_token,501            "eos_token": None,502            "bos_token": None,503            "word_delimiter_token": word_delimiter_token,504        }505 506    # 5. Now we can instantiate the feature extractor, tokenizer and model507    # Note for distributed training, the .from_pretrained methods guarantee that only508    # one local process can concurrently download model & vocab.509 510    # load feature_extractor and tokenizer511    tokenizer = AutoTokenizer.from_pretrained(512        tokenizer_name_or_path,513        use_auth_token=data_args.use_auth_token,514        **tokenizer_kwargs,515    )516    feature_extractor = AutoFeatureExtractor.from_pretrained(517        model_args.model_name_or_path, cache_dir=model_args.cache_dir, use_auth_token=data_args.use_auth_token518    )519 520    # adapt config521    config.update(522        {523            "feat_proj_dropout": model_args.feat_proj_dropout,524            "attention_dropout": model_args.attention_dropout,525            "hidden_dropout": model_args.hidden_dropout,526            "final_dropout": model_args.final_dropout,527            "mask_time_prob": model_args.mask_time_prob,528            "mask_time_length": model_args.mask_time_length,529            "mask_feature_prob": model_args.mask_feature_prob,530            "mask_feature_length": model_args.mask_feature_length,531            "gradient_checkpointing": training_args.gradient_checkpointing,532            "layerdrop": model_args.layerdrop,533            "ctc_loss_reduction": model_args.ctc_loss_reduction,534            "pad_token_id": tokenizer.pad_token_id,535            "vocab_size": len(tokenizer),536            "activation_dropout": model_args.activation_dropout,537        }538    )539 540    # create model541    model = AutoModelForCTC.from_pretrained(542        model_args.model_name_or_path,543        cache_dir=model_args.cache_dir,544        config=config,545        use_auth_token=data_args.use_auth_token,546    )547 548    # freeze encoder549    if model_args.freeze_feature_encoder:550        model.freeze_feature_encoder()551 552    # 6. Now we preprocess the datasets including loading the audio, resampling and normalization553    # Thankfully, `datasets` takes care of automatically loading and resampling the audio,554    # so that we just need to set the correct target sampling rate and normalize the input555    # via the `feature_extractor`556 557    # make sure that dataset decodes audio with correct sampling rate558    dataset_sampling_rate = next(iter(raw_datasets.values())).features[data_args.audio_column_name].sampling_rate559    if dataset_sampling_rate != feature_extractor.sampling_rate:560        raw_datasets = raw_datasets.cast_column(561            data_args.audio_column_name, datasets.features.Audio(sampling_rate=feature_extractor.sampling_rate)562        )563 564    # derive max & min input length for sample rate & max duration565    max_input_length = data_args.max_duration_in_seconds * feature_extractor.sampling_rate566    min_input_length = data_args.min_duration_in_seconds * feature_extractor.sampling_rate567    audio_column_name = data_args.audio_column_name568    num_workers = data_args.preprocessing_num_workers569 570    # `phoneme_language` is only relevant if the model is fine-tuned on phoneme classification571    phoneme_language = data_args.phoneme_language572 573    # Preprocessing the datasets.574    # We need to read the audio files as arrays and tokenize the targets.575    def prepare_dataset(batch):576        # load audio577        sample = batch[audio_column_name]578 579        inputs = feature_extractor(sample["array"], sampling_rate=sample["sampling_rate"])580        batch["input_values"] = inputs.input_values[0]581        batch["input_length"] = len(batch["input_values"])582 583        # encode targets584        additional_kwargs = {}585        if phoneme_language is not None:586            additional_kwargs["phonemizer_lang"] = phoneme_language587 588        batch["labels"] = tokenizer(batch["target_text"], **additional_kwargs).input_ids589        return batch590 591    with training_args.main_process_first(desc="dataset map preprocessing"):592        vectorized_datasets = raw_datasets.map(593            prepare_dataset,594            remove_columns=next(iter(raw_datasets.values())).column_names,595            num_proc=num_workers,596            desc="preprocess datasets",597        )598 599        def is_audio_in_length_range(length):600            return length > min_input_length and length < max_input_length601 602        # filter data that is shorter than min_input_length603        vectorized_datasets = vectorized_datasets.filter(604            is_audio_in_length_range,605            num_proc=num_workers,606            input_columns=["input_length"],607        )608 609    # 7. Next, we can prepare the training.610    # Let's use word error rate (WER) as our evaluation metric,611    # instantiate a data collator and the trainer612 613    # Define evaluation metrics during training, *i.e.* word error rate, character error rate614    eval_metrics = {metric: load_metric(metric) for metric in data_args.eval_metrics}615 616    # for large datasets it is advised to run the preprocessing on a617    # single machine first with ``args.preprocessing_only`` since there will mostly likely618    # be a timeout when running the script in distributed mode.619    # In a second step ``args.preprocessing_only`` can then be set to `False` to load the620    # cached dataset621    if data_args.preprocessing_only:622        logger.info(f"Data preprocessing finished. Files cached at {vectorized_datasets.cache_files}")623        return624 625    def compute_metrics(pred):626        pred_logits = pred.predictions627        pred_ids = np.argmax(pred_logits, axis=-1)628 629        pred.label_ids[pred.label_ids == -100] = tokenizer.pad_token_id630 631        pred_str = tokenizer.batch_decode(pred_ids)632        # we do not want to group tokens when computing the metrics633        label_str = tokenizer.batch_decode(pred.label_ids, group_tokens=False)634 635        metrics = {k: v.compute(predictions=pred_str, references=label_str) for k, v in eval_metrics.items()}636 637        return metrics638 639    # Now save everything to be able to create a single processor later640    if is_main_process(training_args.local_rank):641        # save feature extractor, tokenizer and config642        feature_extractor.save_pretrained(training_args.output_dir)643        tokenizer.save_pretrained(training_args.output_dir)644        config.save_pretrained(training_args.output_dir)645 646    try:647        processor = AutoProcessor.from_pretrained(training_args.output_dir)648    except (OSError, KeyError):649        warnings.warn(650            "Loading a processor from a feature extractor config that does not"651            " include a `processor_class` attribute is deprecated and will be removed in v5. Please add the following "652            " attribute to your `preprocessor_config.json` file to suppress this warning: "653            " `'processor_class': 'Wav2Vec2Processor'`",654            FutureWarning,655        )656        processor = Wav2Vec2Processor.from_pretrained(training_args.output_dir)657 658    # Instantiate custom data collator659    data_collator = DataCollatorCTCWithPadding(processor=processor)660 661    # Initialize Trainer662    trainer = Trainer(663        model=model,664        data_collator=data_collator,665        args=training_args,666        compute_metrics=compute_metrics,667        train_dataset=vectorized_datasets["train"] if training_args.do_train else None,668        eval_dataset=vectorized_datasets["eval"] if training_args.do_eval else None,669        tokenizer=feature_extractor,670    )671 672    # 8. Finally, we can start training673 674    # Training675    if training_args.do_train:676 677        # use last checkpoint if exist678        if last_checkpoint is not None:679            checkpoint = last_checkpoint680        elif os.path.isdir(model_args.model_name_or_path):681            checkpoint = model_args.model_name_or_path682        else:683            checkpoint = None684 685        train_result = trainer.train(resume_from_checkpoint=checkpoint)686        trainer.save_model()687 688        metrics = train_result.metrics689        max_train_samples = (690            data_args.max_train_samples691            if data_args.max_train_samples is not None692            else len(vectorized_datasets["train"])693        )694        metrics["train_samples"] = min(max_train_samples, len(vectorized_datasets["train"]))695 696        trainer.log_metrics("train", metrics)697        trainer.save_metrics("train", metrics)698        trainer.save_state()699 700    # Evaluation701    results = {}702    if training_args.do_eval:703        logger.info("*** Evaluate ***")704        metrics = trainer.evaluate()705        max_eval_samples = (706            data_args.max_eval_samples if data_args.max_eval_samples is not None else len(vectorized_datasets["eval"])707        )708        metrics["eval_samples"] = min(max_eval_samples, len(vectorized_datasets["eval"]))709 710        trainer.log_metrics("eval", metrics)711        trainer.save_metrics("eval", metrics)712 713    # Write model card and (optionally) push to hub714    config_name = data_args.dataset_config_name if data_args.dataset_config_name is not None else "na"715    kwargs = {716        "finetuned_from": model_args.model_name_or_path,717        "tasks": "speech-recognition",718        "tags": ["automatic-speech-recognition", data_args.dataset_name],719        "dataset_args": f"Config: {config_name}, Training split: {data_args.train_split_name}, Eval split: {data_args.eval_split_name}",720        "dataset": f"{data_args.dataset_name.upper()} - {config_name.upper()}",721    }722    if "common_voice" in data_args.dataset_name:723        kwargs["language"] = config_name724 725    if training_args.push_to_hub:726        trainer.push_to_hub(**kwargs)727    else:728        trainer.create_model_card(**kwargs)729 730    return results731 732 733if __name__ == "__main__":734    main()735