NbAiLabArchive/test_w5_long_dataset

1#!/usr/bin/env python32# coding=utf-83# Copyright 2021 The HuggingFace 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"""17Fine-tuning the library models for masked language modeling (BERT, ALBERT, RoBERTa...) with whole word masking on a18text file or a dataset.19 20Here is the full list of checkpoints on the hub that can be fine-tuned by this script:21https://huggingface.co/models?filter=masked-lm22"""23import logging24import os25import sys26import time27from collections import defaultdict28from dataclasses import dataclass, field29 30# You can also adapt this script on your own masked language modeling task. Pointers for this are left as comments.31from pathlib import Path32from typing import Dict, List, Optional, Tuple33 34import datasets35import numpy as np36from datasets import load_dataset37from tqdm import tqdm38 39import flax40import jax41import jax.numpy as jnp42import optax43from flax import jax_utils, traverse_util44from flax.training import train_state45from flax.training.common_utils import get_metrics, onehot, shard46from transformers import (47    CONFIG_MAPPING,48    FLAX_MODEL_FOR_MASKED_LM_MAPPING,49    AutoConfig,50    AutoTokenizer,51    FlaxAutoModelForMaskedLM,52    HfArgumentParser,53    PreTrainedTokenizerBase,54    TensorType,55    TrainingArguments,56    is_tensorboard_available,57    set_seed,58)59 60 61#if datasets.__version__ <= "1.8.0":62#    raise ValueError("Make sure to upgrade `datasets` to a version >= 1.9.0 to use dataset streaming")63 64 65MODEL_CONFIG_CLASSES = list(FLAX_MODEL_FOR_MASKED_LM_MAPPING.keys())66MODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)67 68 69@dataclass70class ModelArguments:71    """72    Arguments pertaining to which model/config/tokenizer we are going to fine-tune, or train from scratch.73    """74 75    model_name_or_path: Optional[str] = field(76        default=None,77        metadata={78            "help": "The model checkpoint for weights initialization."79            "Don't set if you want to train a model from scratch."80        },81    )82    model_type: Optional[str] = field(83        default=None,84        metadata={"help": "If training from scratch, pass a model type from the list: " + ", ".join(MODEL_TYPES)},85    )86    config_name: Optional[str] = field(87        default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}88    )89    tokenizer_name: Optional[str] = field(90        default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}91    )92    cache_dir: Optional[str] = field(93        default=None, metadata={"help": "Where do you want to store the pretrained models downloaded from s3"}94    )95    use_fast_tokenizer: bool = field(96        default=True,97        metadata={"help": "Whether to use one of the fast tokenizer (backed by the tokenizers library) or not."},98    )99    dtype: Optional[str] = field(100        default="float32",101        metadata={102            "help": "Floating-point format in which the model weights should be initialized and trained. Choose one of `[float32, float16, bfloat16]`."103        },104    )105 106 107@dataclass108class DataTrainingArguments:109    """110    Arguments pertaining to what data we are going to input our model for training and eval.111    """112 113    dataset_name: Optional[str] = field(114        default=None, metadata={"help": "The name of the dataset to use (via the datasets library)."}115    )116    dataset_config_name: Optional[str] = field(117        default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}118    )119    train_file: Optional[str] = field(default=None, metadata={"help": "The input training data file (a text file)."})120    validation_file: Optional[str] = field(121        default=None,122        metadata={"help": "An optional input evaluation data file to evaluate the perplexity on (a text file)."},123    )124    train_ref_file: Optional[str] = field(125        default=None,126        metadata={"help": "An optional input train ref data file for whole word masking in Chinese."},127    )128    validation_ref_file: Optional[str] = field(129        default=None,130        metadata={"help": "An optional input validation ref data file for whole word masking in Chinese."},131    )132    auth_token: bool = field(133        default=False, metadata={"help": "Use authorisation token"}134    )135    overwrite_cache: bool = field(136        default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}137    )138    validation_split_percentage: Optional[int] = field(139        default=5,140        metadata={141            "help": "The percentage of the train set used as validation set in case there's no validation split"142        },143    )144    max_seq_length: Optional[int] = field(145        default=None,146        metadata={147            "help": "The maximum total input sequence length after tokenization. Sequences longer "148            "than this will be truncated. Default to the max input length of the model."149        },150    )151    preprocessing_num_workers: Optional[int] = field(152        default=None,153        metadata={"help": "The number of processes to use for the preprocessing."},154    )155    mlm_probability: float = field(156        default=0.15, metadata={"help": "Ratio of tokens to mask for masked language modeling loss"}157    )158    pad_to_max_length: bool = field(159        default=False,160        metadata={161            "help": "Whether to pad all samples to `max_seq_length`. "162            "If False, will pad the samples dynamically when batching to the maximum length in the batch."163        },164    )165    line_by_line: bool = field(166        default=False,167        metadata={"help": "Whether distinct lines of text in the dataset are to be handled as distinct sequences."},168    )169    text_column_name: str = field(170        default="text", metadata={"help": "The name of the column to retrieve the training text."}171    )172    shuffle_buffer_size: int = field(173        default=10000, metadata={"help": "The number of examples to pre-load for shuffling."}174    )175    num_train_steps: int = field(default=50000, metadata={"help": "The number of training steps."})176    num_eval_samples: int = field(default=50000, metadata={"help": "The number of samples to be used for evaluation"})177 178    def __post_init__(self):179        if self.dataset_name is None and self.train_file is None and self.validation_file is None:180            raise ValueError("Need either a dataset name or a training/validation file.")181        else:182            if self.train_file is not None:183                extension = self.train_file.split(".")[-1]184                assert extension in ["csv", "json", "txt"], "`train_file` should be a csv, a json or a txt file."185            if self.validation_file is not None:186                extension = self.validation_file.split(".")[-1]187                assert extension in ["csv", "json", "txt"], "`validation_file` should be a csv, a json or a txt file."188 189 190@flax.struct.dataclass191class FlaxDataCollatorForLanguageModeling:192    """193    Data collator used for language modeling. Inputs are dynamically padded to the maximum length of a batch if they194    are not all of the same length.195 196    Args:197        tokenizer (:class:`~transformers.PreTrainedTokenizer` or :class:`~transformers.PreTrainedTokenizerFast`):198            The tokenizer used for encoding the data.199        mlm_probability (:obj:`float`, `optional`, defaults to 0.15):200            The probability with which to (randomly) mask tokens in the input.201 202    .. note::203 204        For best performance, this data collator should be used with a dataset having items that are dictionaries or205        BatchEncoding, with the :obj:`"special_tokens_mask"` key, as returned by a206        :class:`~transformers.PreTrainedTokenizer` or a :class:`~transformers.PreTrainedTokenizerFast` with the207        argument :obj:`return_special_tokens_mask=True`.208    """209 210    tokenizer: PreTrainedTokenizerBase211    mlm_probability: float = 0.15212 213    def __post_init__(self):214        if self.tokenizer.mask_token is None:215            raise ValueError(216                "This tokenizer does not have a mask token which is necessary for masked language modeling. "217                "You should pass `mlm=False` to train on causal language modeling instead."218            )219 220    def __call__(self, examples: List[Dict[str, np.ndarray]]) -> Dict[str, np.ndarray]:221        # Handle dict or lists with proper padding and conversion to tensor.222        batch = self.tokenizer.pad(examples, return_tensors=TensorType.NUMPY)223 224        # If special token mask has been preprocessed, pop it from the dict.225        special_tokens_mask = batch.pop("special_tokens_mask", None)226 227        batch["input_ids"], batch["labels"] = self.mask_tokens(228            batch["input_ids"], special_tokens_mask=special_tokens_mask229        )230        return batch231 232    def mask_tokens(233        self, inputs: np.ndarray, special_tokens_mask: Optional[np.ndarray]234    ) -> Tuple[jnp.ndarray, jnp.ndarray]:235        """236        Prepare masked tokens inputs/labels for masked language modeling: 80% MASK, 10% random, 10% original.237        """238        labels = inputs.copy()239        # We sample a few tokens in each sequence for MLM training (with probability `self.mlm_probability`)240        probability_matrix = np.full(labels.shape, self.mlm_probability)241        special_tokens_mask = special_tokens_mask.astype("bool")242 243        probability_matrix[special_tokens_mask] = 0.0244        masked_indices = np.random.binomial(1, probability_matrix).astype("bool")245        labels[~masked_indices] = -100  # We only compute loss on masked tokens246 247        # 80% of the time, we replace masked input tokens with tokenizer.mask_token ([MASK])248        indices_replaced = np.random.binomial(1, np.full(labels.shape, 0.8)).astype("bool") & masked_indices249        inputs[indices_replaced] = self.tokenizer.convert_tokens_to_ids(self.tokenizer.mask_token)250 251        # 10% of the time, we replace masked input tokens with random word252        indices_random = np.random.binomial(1, np.full(labels.shape, 0.5)).astype("bool")253        indices_random &= masked_indices & ~indices_replaced254 255        random_words = np.random.randint(self.tokenizer.vocab_size, size=labels.shape, dtype="i4")256        inputs[indices_random] = random_words[indices_random]257 258        # The rest of the time (10% of the time) we keep the masked input tokens unchanged259        return inputs, labels260 261 262def generate_batch_splits(samples_idx: jnp.ndarray, batch_size: int) -> jnp.ndarray:263    num_samples = len(samples_idx)264    samples_to_remove = num_samples % batch_size265 266    if samples_to_remove != 0:267        samples_idx = samples_idx[:-samples_to_remove]268    sections_split = num_samples // batch_size269    batch_idx = np.split(samples_idx, sections_split)270    return batch_idx271 272 273def advance_iter_and_group_samples(train_iterator, num_samples, max_seq_length):274    """275    The training iterator is advanced so that after groupifying the samples,276    `num_samples` of length `max_seq_length` are returned.277    """278    num_total_tokens = max_seq_length * num_samples279    samples = defaultdict(list)280 281    i = 0282    while i < num_total_tokens:283        tokenized_samples = next(train_iterator)284        i += len(tokenized_samples["input_ids"])285 286        # concatenate tokenized samples to list287        samples = {k: samples[k] + tokenized_samples[k] for k in tokenized_samples.keys()}288 289    # Concatenated tokens are split to lists of length `max_seq_length`.290    # Note that remainedr of % max_seq_length are thrown away.291    def group_texts(examples):292        result = {293            k: [t[i : i + max_seq_length] for i in range(0, num_total_tokens, max_seq_length)]294            for k, t in examples.items()295        }296        return result297 298    grouped_samples = group_texts(samples)299    return grouped_samples300 301 302def write_train_metric(summary_writer, train_metrics, train_time, step):303    summary_writer.scalar("train_time", train_time, step)304 305    train_metrics = get_metrics(train_metrics)306    for key, vals in train_metrics.items():307        tag = f"train_{key}"308        for i, val in enumerate(vals):309            summary_writer.scalar(tag, val, step - len(vals) + i + 1)310 311 312def write_eval_metric(summary_writer, eval_metrics, step):313    for metric_name, value in eval_metrics.items():314        summary_writer.scalar(f"eval_{metric_name}", value, step)315 316 317if __name__ == "__main__":318    # See all possible arguments in src/transformers/training_args.py319    # or by passing the --help flag to this script.320    # We now keep distinct sets of args, for a cleaner separation of concerns.321 322    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))323    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):324        # If we pass only one argument to the script and it's the path to a json file,325        # let's parse it to get our arguments.326        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))327    else:328        model_args, data_args, training_args = parser.parse_args_into_dataclasses()329 330    if (331        os.path.exists(training_args.output_dir)332        and os.listdir(training_args.output_dir)333        and training_args.do_train334        and not training_args.overwrite_output_dir335    ):336        raise ValueError(337            f"Output directory ({training_args.output_dir}) already exists and is not empty."338            "Use --overwrite_output_dir to overcome."339        )340 341    # Setup logging342    logging.basicConfig(343        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",344        level="INFO",345        datefmt="[%X]",346    )347 348    # Log on each process the small summary:349    logger = logging.getLogger(__name__)350    #logger.warning(351    #    f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"352    #    + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"353    #)354 355    # Set the verbosity to info of the Transformers logger (on main process only):356    logger.info(f"Training/evaluation parameters {training_args}")357 358    # Set seed before initializing model.359    set_seed(training_args.seed)360 361    # Get the datasets: you can either provide your own CSV/JSON/TXT training and evaluation files (see below)362    # or just provide the name of one of the public datasets available on the hub at https://huggingface.co/datasets/363    # (the dataset will be downloaded automatically from the datasets Hub).364    #365    # For CSV/JSON files, this script will use the column called 'text' or the first column if no column called366    # 'text' is found. You can easily tweak this behavior (see below).367    368    if data_args.dataset_name is not None:369        # Downloading and loading a dataset from the hub.370        dataset = load_dataset(371            data_args.dataset_name,372            data_args.dataset_config_name,373            cache_dir=model_args.cache_dir,374            streaming=True,375            use_auth_token=data_args.auth_token,376            split="train",377        )378 379    if model_args.config_name:380        config = AutoConfig.from_pretrained(model_args.config_name, cache_dir=model_args.cache_dir)381    elif model_args.model_name_or_path:382        config = AutoConfig.from_pretrained(model_args.model_name_or_path, cache_dir=model_args.cache_dir)383    else:384        config = CONFIG_MAPPING[model_args.model_type]()385        logger.warning("You are instantiating a new config instance from scratch.")386 387    if model_args.tokenizer_name:388        tokenizer = AutoTokenizer.from_pretrained(389            model_args.tokenizer_name, cache_dir=model_args.cache_dir, use_fast=model_args.use_fast_tokenizer390        )391    elif model_args.model_name_or_path:392        tokenizer = AutoTokenizer.from_pretrained(393            model_args.model_name_or_path, cache_dir=model_args.cache_dir, use_fast=model_args.use_fast_tokenizer394        )395    else:396        raise ValueError(397            "You are instantiating a new tokenizer from scratch. This is not supported by this script."398            "You can do it from another script, save it, and load it from here, using --tokenizer_name."399        )400 401    # Otherwise, we tokenize every text, then concatenate them together before splitting them in smaller parts.402    # We use `return_special_tokens_mask=True` because DataCollatorForLanguageModeling (see below) is more403    # efficient when it receives the `special_tokens_mask`.404    def tokenize_function(examples):405        return tokenizer(examples[data_args.text_column_name], return_special_tokens_mask=True)406 407    tokenized_datasets = dataset.map(408        tokenize_function,409        batched=True,410    )411 412    shuffle_seed = training_args.seed413    tokenized_datasets = tokenized_datasets.shuffle(buffer_size=data_args.shuffle_buffer_size, seed=shuffle_seed)414 415    has_tensorboard = is_tensorboard_available()416    if has_tensorboard and jax.process_index() == 0:417        try:418            from flax.metrics.tensorboard import SummaryWriter419        except ImportError as ie:420            has_tensorboard = False421            logger.warning(422                f"Unable to display metrics through TensorBoard because some package are not installed: {ie}"423            )424 425        summary_writer = SummaryWriter(log_dir=Path(training_args.output_dir))426 427    # Data collator428    # This one will take care of randomly masking the tokens.429    data_collator = FlaxDataCollatorForLanguageModeling(tokenizer=tokenizer, mlm_probability=data_args.mlm_probability)430 431    # Initialize our training432    rng = jax.random.PRNGKey(training_args.seed)433    dropout_rngs = jax.random.split(rng, jax.local_device_count())434 435    if model_args.model_name_or_path:436        model = FlaxAutoModelForMaskedLM.from_pretrained(437            model_args.model_name_or_path, config=config, seed=training_args.seed, dtype=getattr(jnp, model_args.dtype)438        )439    else:440        model = FlaxAutoModelForMaskedLM.from_config(441            config, seed=training_args.seed, dtype=getattr(jnp, model_args.dtype)442        )443 444    # Store some constant445    num_epochs = int(training_args.num_train_epochs)446    train_batch_size = int(training_args.per_device_train_batch_size) * jax.device_count()447    eval_batch_size = int(training_args.per_device_eval_batch_size) * jax.device_count()448 449    # define number steps per stream epoch450    num_train_steps = data_args.num_train_steps451 452    # Create learning rate schedule453    warmup_fn = optax.linear_schedule(454        init_value=0.0, end_value=training_args.learning_rate, transition_steps=training_args.warmup_steps455    )456    decay_fn = optax.linear_schedule(457        init_value=training_args.learning_rate,458        end_value=0,459        transition_steps=num_train_steps - training_args.warmup_steps,460    )461    linear_decay_lr_schedule_fn = optax.join_schedules(462        schedules=[warmup_fn, decay_fn], boundaries=[training_args.warmup_steps]463    )464 465    # We use Optax's "masking" functionality to not apply weight decay466    # to bias and LayerNorm scale parameters. decay_mask_fn returns a467    # mask boolean with the same structure as the parameters.468    # The mask is True for parameters that should be decayed.469    # Note that this mask is specifically adapted for FlaxBERT-like models.470    # For other models, one should correct the layer norm parameter naming471    # accordingly.472    def decay_mask_fn(params):473        flat_params = traverse_util.flatten_dict(params)474        flat_mask = {path: (path[-1] != "bias" and path[-2:] != ("LayerNorm", "scale")) for path in flat_params}475        return traverse_util.unflatten_dict(flat_mask)476 477 478 479    #te adam optimizer480    if training_args.adafactor:481        # We use the default parameters here to initialize adafactor,482        # For more details about the parameters please check https://github.com/deepmind/optax/blob/ed02befef9bf81cbbf236be3d2b0e032e9ed4a40/optax/_src/alias.py#L74483        optimizer = optax.adafactor(484            learning_rate=linear_decay_lr_schedule_fn,485        )486    else:487        optimizer = optax.adamw(488            learning_rate=linear_decay_lr_schedule_fn,489            b1=training_args.adam_beta1,490            b2=training_args.adam_beta2,491            eps=training_args.adam_epsilon,492            weight_decay=training_args.weight_decay,493            mask=decay_mask_fn,494        )495 496    # Setup train state497    state = train_state.TrainState.create(apply_fn=model.__call__, params=model.params, tx=optimizer)498 499    # Define gradient update step fn500    def train_step(state, batch, dropout_rng):501        dropout_rng, new_dropout_rng = jax.random.split(dropout_rng)502 503        def loss_fn(params):504            labels = batch.pop("labels")505 506            logits = state.apply_fn(**batch, params=params, dropout_rng=dropout_rng, train=True)[0]507 508            # compute loss, ignore padded input tokens509            label_mask = jnp.where(labels > 0, 1.0, 0.0)510            loss = optax.softmax_cross_entropy(logits, onehot(labels, logits.shape[-1])) * label_mask511 512            # take average513            loss = loss.sum() / label_mask.sum()514 515            return loss516 517        grad_fn = jax.value_and_grad(loss_fn)518        loss, grad = grad_fn(state.params)519        grad = jax.lax.pmean(grad, "batch")520        new_state = state.apply_gradients(grads=grad)521 522        metrics = jax.lax.pmean(523            {"loss": loss, "learning_rate": linear_decay_lr_schedule_fn(state.step)}, axis_name="batch"524        )525 526        return new_state, metrics, new_dropout_rng527 528    # Create parallel version of the train step529    p_train_step = jax.pmap(train_step, "batch", donate_argnums=(0,))530 531    # Define eval fn532    def eval_step(params, batch):533        labels = batch.pop("labels")534 535        logits = model(**batch, params=params, train=False)[0]536 537        # compute loss, ignore padded input tokens538        label_mask = jnp.where(labels > 0, 1.0, 0.0)539        loss = optax.softmax_cross_entropy(logits, onehot(labels, logits.shape[-1])) * label_mask540 541        # compute accuracy542        accuracy = jnp.equal(jnp.argmax(logits, axis=-1), labels) * label_mask543 544        # summarize metrics545        metrics = {"loss": loss.sum(), "accuracy": accuracy.sum(), "normalizer": label_mask.sum()}546        metrics = jax.lax.psum(metrics, axis_name="batch")547 548        return metrics549 550    p_eval_step = jax.pmap(eval_step, "batch", donate_argnums=(0,))551 552    # Replicate the train state on each device553    state = jax_utils.replicate(state)554 555    train_time = 0556    train_start = time.time()557    train_metrics = []558    eval_metrics = []559 560    training_iter = iter(tokenized_datasets)561 562    max_seq_length = min(data_args.max_seq_length, tokenizer.model_max_length)563    eval_samples = advance_iter_and_group_samples(training_iter, data_args.num_eval_samples, max_seq_length)564 565    steps = tqdm(range(num_train_steps), desc="Training...", position=0)566    for step in range(num_train_steps):567        # ======================== Training ================================568        try:569            samples = advance_iter_and_group_samples(training_iter, train_batch_size, max_seq_length)570        except StopIteration:571            # Once the end of the dataset stream is reached, the training iterator572            # is reinitialized and reshuffled and a new eval dataset is randomely chosen.573            shuffle_seed += 1574            tokenized_datasets.set_epoch(shuffle_seed)575 576            training_iter = iter(tokenized_datasets)577 578            eval_dataset = advance_iter_and_group_samples(training_iter, data_args.num_eval_samples, max_seq_length)579            samples = advance_iter_and_group_samples(training_iter, train_batch_size, max_seq_length)580 581        # process input samples582        model_inputs = data_collator(samples)583 584        # Model forward585        model_inputs = shard(model_inputs.data)586        state, train_metric, dropout_rngs = p_train_step(state, model_inputs, dropout_rngs)587 588        train_metrics.append(train_metric)589 590        if step % training_args.logging_steps == 0 and step > 0:591            steps.write(592                f"Step... ({step} | Loss: {train_metric['loss'].mean()}, Learning Rate: {train_metric['learning_rate'].mean()})"593            )594            train_time += time.time() - train_start595            if has_tensorboard and jax.process_index() == 0:596                write_train_metric(summary_writer, train_metrics, train_time, step)597            train_metrics = []598 599        # ======================== Evaluating ==============================600        if step % training_args.eval_steps == 0 and step > 0:601            eval_samples_idx = jnp.arange(data_args.num_eval_samples)602            eval_batch_idx = generate_batch_splits(eval_samples_idx, eval_batch_size)603 604            for i, batch_idx in enumerate(tqdm(eval_batch_idx, desc="Evaluating ...", position=1)):605                # process input samples606                batch_eval_samples = {k: [v[idx] for idx in batch_idx] for k, v in eval_samples.items()}607                model_inputs = data_collator(batch_eval_samples)608 609                # Model forward610                model_inputs = shard(model_inputs.data)611                metrics = p_eval_step(state.params, model_inputs)612                eval_metrics.append(metrics)613 614            # normalize eval metrics615            eval_metrics = get_metrics(eval_metrics)616            eval_metrics = jax.tree_map(jnp.sum, eval_metrics)617            eval_normalizer = eval_metrics.pop("normalizer")618            eval_metrics = jax.tree_map(lambda x: x / eval_normalizer, eval_metrics)619 620            # Update progress bar621            steps.desc = f"Step... ({step + 1}/{num_train_steps} | Loss: {eval_metrics['loss']}, Acc: {eval_metrics['accuracy']})"622 623            if has_tensorboard and jax.process_index() == 0:624                write_eval_metric(summary_writer, eval_metrics, step)625            eval_metrics = []626 627            # save checkpoint after each epoch and push checkpoint to the hub628            if jax.process_index() == 0:629                params = jax.device_get(jax.tree_map(lambda x: x[0], state.params))630                model.save_pretrained(631                    training_args.output_dir,632                    params=params,633                    push_to_hub=training_args.push_to_hub,634                    commit_message=f"Saving weights and logs of step {step+1}",635                )636 637        # update tqdm bar638        steps.update(1)639