


"""simple docstring"""



import argparse
import logging
import os
import re

import tensorflow as tf

from transformers import (
    AutoConfig,
    AutoTokenizer,
    DataCollatorForLanguageModeling,
    PushToHubCallback,
    TFAutoModelForMaskedLM,
    create_optimizer,
)


__snake_case       =       logging.getLogger(__name__)

__snake_case       =       tf.data.AUTOTUNE







def     __lowerCAmelCase		(		)      ->		Dict:


    """simple docstring"""





    snake_case       : Tuple    =       argparse.ArgumentParser(description="Train a masked language model on TPU."					)
    parser.add_argument(
        "--pretrained_model_config"		,				type=lowercase		,				default="roberta-base"		,				help="The model config to use. Note that we don't copy the model's weights, only the config!"		,				)
    parser.add_argument(
        "--tokenizer"		,				type=lowercase		,				default="unigram-tokenizer-wikitext"		,				help="The name of the tokenizer to load. We use the pretrained tokenizer to initialize the model's vocab size."		,				)

    parser.add_argument(
        "--per_replica_batch_size"		,				type=lowercase		,				default=8		,				help="Batch size per TPU core."		,				)

    parser.add_argument(
        "--no_tpu"		,				action="store_true"		,				help="If set, run on CPU and don't try to initialize a TPU. Useful for debugging on non-TPU instances."		,				)

    parser.add_argument(
        "--tpu_name"		,				type=lowercase		,				help="Name of TPU resource to initialize. Should be blank on Colab, and 'local' on TPU VMs."		,				default="local"		,				)

    parser.add_argument(
        "--tpu_zone"		,				type=lowercase		,				help="Google cloud zone that TPU resource is located in. Only used for non-Colab TPU nodes."		,				)

    parser.add_argument(
        "--gcp_project"		,				type=lowercase		,				help="Google cloud project name. Only used for non-Colab TPU nodes."					)

    parser.add_argument(
        "--bfloat16"		,				action="store_true"		,				help="Use mixed-precision bfloat16 for training. This is the recommended lower-precision format for TPU."		,				)

    parser.add_argument(
        "--train_dataset"		,				type=lowercase		,				help="Path to training dataset to load. If the path begins with `gs://`"
        " then the dataset will be loaded from a Google Cloud Storage bucket."		,				)

    parser.add_argument(
        "--shuffle_buffer_size"		,				type=lowercase		,				default=2**18		,				help="Size of the shuffle buffer (in samples)"		,				)

    parser.add_argument(
        "--eval_dataset"		,				type=lowercase		,				help="Path to evaluation dataset to load. If the path begins with `gs://`"
        " then the dataset will be loaded from a Google Cloud Storage bucket."		,				)

    parser.add_argument(
        "--num_epochs"		,				type=lowercase		,				default=1		,				help="Number of epochs to train for."		,				)

    parser.add_argument(
        "--learning_rate"		,				type=lowercase		,				default=1e-4		,				help="Learning rate to use for training."		,				)

    parser.add_argument(
        "--weight_decay_rate"		,				type=lowercase		,				default=1e-3		,				help="Weight decay rate to use for training."		,				)

    parser.add_argument(
        "--max_length"		,				type=lowercase		,				default=512		,				help="Maximum length of tokenized sequences. Should match the setting used in prepare_tfrecord_shards.py"		,				)

    parser.add_argument(
        "--mlm_probability"		,				type=lowercase		,				default=0.15		,				help="Fraction of tokens to mask during training."		,				)

    parser.add_argument("--output_dir"		,				type=lowercase		,				required=lowercase		,				help="Path to save model checkpoints to."					)
    parser.add_argument("--hub_model_id"		,				type=lowercase		,				help="Model ID to upload to on the Hugging Face Hub."					)

    snake_case       : int    =       parser.parse_args()
    return args







def     __lowerCAmelCase		(		lowercase    :    Dict					)      ->		Tuple:


    """simple docstring"""





    try:
        if args.tpu_name:
            snake_case       : List[Any]    =       tf.distribute.cluster_resolver.TPUClusterResolver(
                args.tpu_name		,				zone=args.tpu_zone		,				project=args.gcp_project					)
        else:
            snake_case       : Optional[int]    =       tf.distribute.cluster_resolver.TPUClusterResolver()
    except ValueError:
        raise RuntimeError(
            "Couldn't connect to TPU! Most likely you need to specify --tpu_name, --tpu_zone, or "
            "--gcp_project. When running on a TPU VM, use --tpu_name local."					)

    tf.config.experimental_connect_to_cluster(lowercase					)
    tf.tpu.experimental.initialize_tpu_system(lowercase					)

    return tpu







def     __lowerCAmelCase		(		lowercase    :    Dict					)      ->		Optional[Any]:


    """simple docstring"""





    snake_case       : Union[str, Any]    =       0
    for file in file_list:
        snake_case       : List[Any]    =       file.split("/"					)[-1]
        snake_case       : Union[str, Any]    =       re.search(R"-\d+-(\d+)\.tfrecord"		,				lowercase					).group(1					)
        snake_case       : Optional[int]    =       int(lowercase					)
        num_samples += sample_count

    return num_samples







def     __lowerCAmelCase		(		lowercase    :    Optional[Any]		,				lowercase    :    Optional[int]		,				lowercase    :    Optional[int]		,				lowercase    :    int		,				lowercase    :    Optional[Any]		,				lowercase    :    Tuple=None					)      ->		str:


    """simple docstring"""





    snake_case       : List[str]    =       count_samples(lowercase					)
    snake_case       : Tuple    =       tf.data.Dataset.from_tensor_slices(lowercase					)
    if shuffle:
        snake_case       : Optional[int]    =       dataset.shuffle(len(lowercase					)					)
    snake_case       : List[Any]    =       tf.data.TFRecordDataset(lowercase		,				num_parallel_reads=lowercase					)
    # TF can't infer the total sample count because it doesn't read all the records yet, so we assert it here
    snake_case       : str    =       dataset.apply(tf.data.experimental.assert_cardinality(lowercase					)					)
    snake_case       : str    =       dataset.map(lowercase		,				num_parallel_calls=lowercase					)
    if shuffle:
        assert shuffle_buffer_size is not None
        snake_case       : List[str]    =       dataset.shuffle(args.shuffle_buffer_size					)
    snake_case       : str    =       dataset.batch(lowercase		,				drop_remainder=lowercase					)
    snake_case       : str    =       dataset.map(lowercase		,				num_parallel_calls=lowercase					)
    snake_case       : Tuple    =       dataset.prefetch(lowercase					)
    return dataset







def     __lowerCAmelCase		(		lowercase    :    List[str]					)      ->		Tuple:


    """simple docstring"""





    if not args.no_tpu:
        snake_case       : Tuple    =       initialize_tpu(lowercase					)
        snake_case       : List[str]    =       tf.distribute.TPUStrategy(lowercase					)
    else:
        snake_case       : Optional[int]    =       tf.distribute.OneDeviceStrategy(device="/gpu:0"					)

    if args.bfloataa:
        tf.keras.mixed_precision.set_global_policy("mixed_bfloat16"					)

    snake_case       : str    =       AutoTokenizer.from_pretrained(args.tokenizer					)
    snake_case       : Optional[Any]    =       AutoConfig.from_pretrained(args.pretrained_model_config					)
    snake_case       : Optional[int]    =       tokenizer.vocab_size

    snake_case       : str    =       tf.io.gfile.glob(os.path.join(args.train_dataset		,				"*.tfrecord"					)					)
    if not training_records:
        raise ValueError(F'No .tfrecord files found in {args.train_dataset}.'					)
    snake_case       : Optional[int]    =       tf.io.gfile.glob(os.path.join(args.eval_dataset		,				"*.tfrecord"					)					)
    if not eval_records:
        raise ValueError(F'No .tfrecord files found in {args.eval_dataset}.'					)

    snake_case       : Union[str, Any]    =       count_samples(lowercase					)

    snake_case       : Dict    =       num_train_samples // (args.per_replica_batch_size * strategy.num_replicas_in_sync)
    snake_case       : int    =       steps_per_epoch * args.num_epochs

    with strategy.scope():
        snake_case       : Dict    =       TFAutoModelForMaskedLM.from_config(lowercase					)
        model(model.dummy_inputs					)  # Pass some dummy inputs through the model to ensure all the weights are built
        snake_case   ,snake_case       : List[str]    =       create_optimizer(
            num_train_steps=lowercase		,				num_warmup_steps=total_train_steps // 20		,				init_lr=args.learning_rate		,				weight_decay_rate=args.weight_decay_rate		,				)

        # Transformers models compute the right loss for their task by default when labels are passed, and will
        # use this for training unless you specify your own loss function in compile().
        model.compile(optimizer=lowercase		,				metrics=["accuracy"]					)

    def decode_fn(lowercase    :    Tuple					):
        snake_case       : Tuple    =       {
            "input_ids": tf.io.FixedLenFeature(dtype=tf.intaa		,				shape=(args.max_length,)					),
            "attention_mask": tf.io.FixedLenFeature(dtype=tf.intaa		,				shape=(args.max_length,)					),
        }
        return tf.io.parse_single_example(lowercase		,				lowercase					)

    # Many of the data collators in Transformers are TF-compilable when return_tensors == "tf", so we can
    # use their methods in our data pipeline.
    snake_case       : List[str]    =       DataCollatorForLanguageModeling(
        tokenizer=lowercase		,				mlm_probability=args.mlm_probability		,				mlm=lowercase		,				return_tensors="tf"					)

    def mask_with_collator(lowercase    :    List[str]					):
        # TF really needs an isin() function
        snake_case       : Any    =       (
            ~tf.cast(batch["attention_mask"]		,				tf.bool					)
            | (batch["input_ids"] == tokenizer.cls_token_id)
            | (batch["input_ids"] == tokenizer.sep_token_id)
        )
        snake_case   ,snake_case       : Optional[Any]    =       data_collator.tf_mask_tokens(
            batch["input_ids"]		,				vocab_size=len(lowercase					)		,				mask_token_id=tokenizer.mask_token_id		,				special_tokens_mask=lowercase		,				)
        return batch

    snake_case       : Optional[Any]    =       args.per_replica_batch_size * strategy.num_replicas_in_sync

    snake_case       : Dict    =       prepare_dataset(
        lowercase		,				decode_fn=lowercase		,				mask_fn=lowercase		,				batch_size=lowercase		,				shuffle=lowercase		,				shuffle_buffer_size=args.shuffle_buffer_size		,				)

    snake_case       : Tuple    =       prepare_dataset(
        lowercase		,				decode_fn=lowercase		,				mask_fn=lowercase		,				batch_size=lowercase		,				shuffle=lowercase		,				)

    snake_case       : List[Any]    =       []
    if args.hub_model_id:
        callbacks.append(
            PushToHubCallback(output_dir=args.output_dir		,				hub_model_id=args.hub_model_id		,				tokenizer=lowercase					)					)

    model.fit(
        lowercase		,				validation_data=lowercase		,				epochs=args.num_epochs		,				callbacks=lowercase		,				)

    model.save_pretrained(args.output_dir					)


if __name__ == "__main__":
  __snake_case       =       parse_args()
  main(args)









"""simple docstring"""



# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from ..models.whisper import WhisperForConditionalGeneration, WhisperProcessor
from .base import PipelineTool







class _lowerCAmelCase							( snake_case_      ):
   __UpperCAmelCase					:			List[str]				    =    '''openai/whisper-base'''
   __UpperCAmelCase					:			Any				    =    (
       '''This is a tool that transcribes an audio into text. It takes an input named `audio` and returns the '''
       '''transcribed text.'''
   )
   __UpperCAmelCase					:			Tuple				    =    '''transcriber'''
   __UpperCAmelCase					:			Union[str, Any]				    =    WhisperProcessor
   __UpperCAmelCase					:			str				    =    WhisperForConditionalGeneration

   __UpperCAmelCase					:			int				    =    ['''audio''']
   __UpperCAmelCase					:			Optional[int]				    =    ['''text''']







   def 					lowerCamelCase						(      self		,      UpperCamelCase__		)      ->						Tuple:

       '''simple docstring'''




       return self.pre_processor(UpperCamelCase__		,      return_tensors="pt"		).input_features







   def 					lowerCamelCase						(      self		,      UpperCamelCase__		)      ->						int:

       '''simple docstring'''




       return self.model.generate(inputs=UpperCamelCase__		)







   def 					lowerCamelCase						(      self		,      UpperCamelCase__		)      ->						Optional[Any]:

       '''simple docstring'''




       return self.pre_processor.batch_decode(UpperCamelCase__		,      skip_special_tokens=UpperCamelCase__		)[0]





