



import json
import logging
import os
import re
import sys
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Union

import datasets
import numpy as np
import torch
import torchaudio
from packaging import version
from torch import nn

import transformers
from transformers import (
    HfArgumentParser,
    Trainer,
    TrainingArguments,
    WavaVecaCTCTokenizer,
    WavaVecaFeatureExtractor,
    WavaVecaForCTC,
    WavaVecaProcessor,
    is_apex_available,
    set_seed,
)
from transformers.trainer_utils import get_last_checkpoint, is_main_process


if is_apex_available():
    from apex import amp


if version.parse(version.parse(torch.__version__).base_version) >= version.parse('''1.6'''):
    __UpperCAmelCase			  =			True
    from torch.cuda.amp import autocast

__UpperCAmelCase			  =			logging.getLogger(__name__)
def 						UpperCamelCase      (  snake_case__			:       Tuple=None			,	snake_case__			:       Dict=None      )      ->	Any:
   return field(default_factory=lambda: default			,	metadata=lowercase_      )



@dataclass
class 							lowerCAmelCase_						:
       UpperCAmelCase__							:   str  	=    field(
           metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}				)
       UpperCAmelCase__							:   Optional[str]  	=    field(
           default=__A			,       metadata={"help": "Where do you want to store the pretrained models downloaded from huggingface.co"}			,       )
       UpperCAmelCase__							:   Optional[bool]  	=    field(
           default=__A			,       metadata={"help": "Whether to freeze the feature extractor layers of the model."}				)
       UpperCAmelCase__							:   Optional[float]  	=    field(
           default=0.1			,       metadata={"help": "The dropout ratio for the attention probabilities."}				)
       UpperCAmelCase__							:   Optional[float]  	=    field(
           default=0.1			,       metadata={"help": "The dropout ratio for activations inside the fully connected layer."}				)
       UpperCAmelCase__							:   Optional[float]  	=    field(
           default=0.1			,       metadata={
               "help": "The dropout probabilitiy for all fully connected layers in the embeddings, encoder, and pooler."
           }			,       )
       UpperCAmelCase__							:   Optional[float]  	=    field(
           default=0.1			,       metadata={"help": "The dropout probabilitiy for all 1D convolutional layers in feature extractor."}			,       )
       UpperCAmelCase__							:   Optional[float]  	=    field(
           default=0.05			,       metadata={
               "help": (
                   "Propability of each feature vector along the time axis to be chosen as the start of the vector"
                   "span to be masked. Approximately ``mask_time_prob * sequence_length // mask_time_length`` feature"
                   "vectors will be masked along the time axis. This is only relevant if ``apply_spec_augment is True``."
               )
           }			,       )
       UpperCAmelCase__							:   Optional[float]  	=    field(default=0.0			,       metadata={"help": "The LayerDrop probability."}				)






@dataclass
class 							lowerCAmelCase_						:
       UpperCAmelCase__							:   Optional[str]  	=    field(
           default=__A			,       metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}				)
       UpperCAmelCase__							:   Optional[str]  	=    field(
           default="train+validation"			,       metadata={
               "help": "The name of the training data set split to use (via the datasets library). Defaults to \'train\'"
           }			,       )
       UpperCAmelCase__							:   bool  	=    field(
           default=__A			,       metadata={"help": "Overwrite the cached preprocessed datasets or not."}				)
       UpperCAmelCase__							:   Optional[int]  	=    field(
           default=__A			,       metadata={"help": "The number of processes to use for the preprocessing."}			,       )
       UpperCAmelCase__							:   Optional[int]  	=    field(
           default=__A			,       metadata={
               "help": (
                   "For debugging purposes or quicker training, truncate the number of training examples to this "
                   "value if set."
               )
           }			,       )
       UpperCAmelCase__							:   Optional[int]  	=    field(
           default=__A			,       metadata={
               "help": (
                   "For debugging purposes or quicker training, truncate the number of validation examples to this "
                   "value if set."
               )
           }			,       )
       UpperCAmelCase__							:   List[str]  	=    list_field(
           default=[",", "?", ".", "!", "-", ";", ":", "\"\"", "%", "\'", "\"", "�"]			,       metadata={"help": "A list of characters to remove from the transcripts."}			,       )






@dataclass
class 							lowerCAmelCase_						:
       UpperCAmelCase__							:   WavaVecaProcessor
       UpperCAmelCase__							:   Union[bool, str]  	=    True
       UpperCAmelCase__							:   Optional[int]  	=    None
       UpperCAmelCase__							:   Optional[int]  	=    None
       UpperCAmelCase__							:   Optional[int]  	=    None
       UpperCAmelCase__							:   Optional[int]  	=    None
       def __call__(     self,     SCREAMING_SNAKE_CASE_			)	->						Any:
          UpperCamelCase       :    Optional[Any]		      =      [{'input_values': feature['input_values']} for feature in features]
          UpperCamelCase       :    List[str]		      =      [{'input_ids': feature['labels']} for feature in features]

          UpperCamelCase       :    Union[str, Any]		      =      self.processor.pad(
              UpperCamelCase__,     padding=self.padding,     max_length=self.max_length,     pad_to_multiple_of=self.pad_to_multiple_of,     return_tensors='pt',     )
          UpperCamelCase       :    Optional[int]		      =      self.processor.pad(
              labels=UpperCamelCase__,     padding=self.padding,     max_length=self.max_length_labels,     pad_to_multiple_of=self.pad_to_multiple_of_labels,     return_tensors='pt',     )

          # replace padding with -100 to ignore loss correctly
          UpperCamelCase       :    str		      =      labels_batch['input_ids'].masked_fill(labels_batch.attention_mask.ne(1			),     -100			)

          UpperCamelCase       :    Optional[Any]		      =      labels

          return batch






class 							lowerCAmelCase_						(    __A				):
       def    snake_case_			(     self,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_			)	->						List[Any]:
          model.train()
          UpperCamelCase       :    Tuple		      =      self._prepare_inputs(UpperCamelCase__			)

          if self.use_amp:
             with autocast():
                UpperCamelCase       :    Tuple		      =      self.compute_loss(UpperCamelCase__,     UpperCamelCase__			)
          else:
             UpperCamelCase       :    int		      =      self.compute_loss(UpperCamelCase__,     UpperCamelCase__			)

          if self.args.n_gpu > 1:
             if model.module.config.ctc_loss_reduction == "mean":
                UpperCamelCase       :    Tuple		      =      loss.mean()
             elif model.module.config.ctc_loss_reduction == "sum":
                UpperCamelCase       :    Any		      =      loss.sum() / (inputs['labels'] >= 0).sum()
             else:
                raise ValueError(F"""{model.config.ctc_loss_reduction} is not valid. Choose one of [\'mean\', \'sum\']"""			)

          if self.args.gradient_accumulation_steps > 1:
             UpperCamelCase       :    int		      =      loss / self.args.gradient_accumulation_steps

          if self.use_amp:
             self.scaler.scale(UpperCamelCase__			).backward()
          elif self.use_apex:
             with amp.scale_loss(UpperCamelCase__,     self.optimizer			) as scaled_loss:
                scaled_loss.backward()
          elif self.deepspeed:
             self.deepspeed.backward(UpperCamelCase__			)
          else:
             loss.backward()

          return loss.detach()






def 						UpperCamelCase      (  )      ->	Union[str, Any]:
   UpperCamelCase       :    Dict		      =      HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)      )
   if len(sys.argv      ) == 2 and sys.argv[1].endswith('.json'      ):
      # If we pass only one argument to the script and it's the path to a json file,
      # let's parse it to get our arguments.
      UpperCamelCase      ,	UpperCamelCase      ,	UpperCamelCase       :    Dict		      =      parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]      )      )
   else:
      UpperCamelCase      ,	UpperCamelCase      ,	UpperCamelCase       :    Dict		      =      parser.parse_args_into_dataclasses()

   # Detecting last checkpoint.
   UpperCamelCase       :    Union[str, Any]		      =      None
   if os.path.isdir(training_args.output_dir      ) and training_args.do_train and not training_args.overwrite_output_dir:
      UpperCamelCase       :    str		      =      get_last_checkpoint(training_args.output_dir      )
      if last_checkpoint is None and len(os.listdir(training_args.output_dir      )      ) > 0:
         raise ValueError(
             F"""Output directory ({training_args.output_dir}) already exists and is not empty. """
             'Use --overwrite_output_dir to overcome.'      )
      elif last_checkpoint is not None:
         logger.info(
             F"""Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change """
             'the `--output_dir` or add `--overwrite_output_dir` to train from scratch.'      )

    # Setup logging
   logging.basicConfig(
       format='%(asctime)s - %(levelname)s - %(name)s - %(message)s'			,	datefmt='%m/%d/%Y %H:%M:%S'			,	handlers=[logging.StreamHandler(sys.stdout      )]			,	)
   logger.setLevel(logging.INFO if is_main_process(training_args.local_rank      ) else logging.WARN      )

   # Log on each process the small summary:
   logger.warning(
       F"""Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"""
       + F"""distributed training: {bool(training_args.local_rank != -1      )}, 16-bits training: {training_args.fpaa}"""      )
   # Set the verbosity to info of the Transformers logger (on main process only):
   if is_main_process(training_args.local_rank      ):
      transformers.utils.logging.set_verbosity_info()
   logger.info('Training/evaluation parameters %s'			,	lowercase_      )

   # Set seed before initializing model.
   set_seed(training_args.seed      )

   # Get the datasets:
   UpperCamelCase       :    Optional[int]		      =      datasets.load_dataset(
       'common_voice'			,	data_args.dataset_config_name			,	split=data_args.train_split_name      )
   UpperCamelCase       :    Any		      =      datasets.load_dataset('common_voice'			,	data_args.dataset_config_name			,	split='test'      )

   # Create and save tokenizer
   UpperCamelCase       :    Optional[int]		      =      F"""[{"".join(data_args.chars_to_ignore      )}]"""

   def remove_special_characters(snake_case__			:       Union[str, Any]      ):
      UpperCamelCase       :    Optional[Any]		      =      re.sub(lowercase_			,	''			,	batch['sentence']      ).lower() + ' '
      return batch

   UpperCamelCase       :    Optional[Any]		      =      train_dataset.map(lowercase_			,	remove_columns=['sentence']      )
   UpperCamelCase       :    Any		      =      eval_dataset.map(lowercase_			,	remove_columns=['sentence']      )

   def extract_all_chars(snake_case__			:       Optional[int]      ):
      UpperCamelCase       :    Dict		      =      ' '.join(batch['text']      )
      UpperCamelCase       :    Union[str, Any]		      =      list(set(lowercase_      )      )
      return {"vocab": [vocab], "all_text": [all_text]}

   UpperCamelCase       :    Optional[int]		      =      train_dataset.map(
       lowercase_			,	batched=lowercase_			,	batch_size=-1			,	keep_in_memory=lowercase_			,	remove_columns=train_dataset.column_names			,	)
   UpperCamelCase       :    str		      =      train_dataset.map(
       lowercase_			,	batched=lowercase_			,	batch_size=-1			,	keep_in_memory=lowercase_			,	remove_columns=eval_dataset.column_names			,	)

   UpperCamelCase       :    Union[str, Any]		      =      list(set(vocab_train['vocab'][0]      ) | set(vocab_test['vocab'][0]      )      )
   UpperCamelCase       :    Optional[Any]		      =      {v: k for k, v in enumerate(lowercase_      )}
   UpperCamelCase       :    Optional[Any]		      =      vocab_dict[' ']
   del vocab_dict[" "]
   UpperCamelCase       :    Optional[Any]		      =      len(lowercase_      )
   UpperCamelCase       :    Union[str, Any]		      =      len(lowercase_      )

   with open('vocab.json'			,	'w'      ) as vocab_file:
      json.dump(lowercase_			,	lowercase_      )

   # Load pretrained model and tokenizer
   #
   # Distributed training:
   # The .from_pretrained methods guarantee that only one local process can concurrently
   # download model & vocab.
   UpperCamelCase       :    List[Any]		      =      WavaVecaCTCTokenizer(
       'vocab.json'			,	unk_token='[UNK]'			,	pad_token='[PAD]'			,	word_delimiter_token='|'			,	)
   UpperCamelCase       :    str		      =      WavaVecaFeatureExtractor(
       feature_size=1			,	sampling_rate=16000			,	padding_value=0.0			,	do_normalize=lowercase_			,	return_attention_mask=lowercase_      )
   UpperCamelCase       :    List[Any]		      =      WavaVecaProcessor(feature_extractor=lowercase_			,	tokenizer=lowercase_      )
   UpperCamelCase       :    str		      =      WavaVecaForCTC.from_pretrained(
       model_args.model_name_or_path			,	cache_dir=model_args.cache_dir			,	activation_dropout=model_args.activation_dropout			,	attention_dropout=model_args.attention_dropout			,	hidden_dropout=model_args.hidden_dropout			,	feat_proj_dropout=model_args.feat_proj_dropout			,	mask_time_prob=model_args.mask_time_prob			,	gradient_checkpointing=training_args.gradient_checkpointing			,	layerdrop=model_args.layerdrop			,	ctc_loss_reduction='mean'			,	pad_token_id=processor.tokenizer.pad_token_id			,	vocab_size=len(processor.tokenizer      )			,	)

   if data_args.max_train_samples is not None:
      UpperCamelCase       :    List[str]		      =      min(len(lowercase_      )			,	data_args.max_train_samples      )
      UpperCamelCase       :    Union[str, Any]		      =      train_dataset.select(range(lowercase_      )      )

   if data_args.max_val_samples is not None:
      UpperCamelCase       :    Any		      =      eval_dataset.select(range(data_args.max_val_samples      )      )

   UpperCamelCase       :    Dict		      =      torchaudio.transforms.Resample(48000			,	16000      )

   # Preprocessing the datasets.
   # We need to read the aduio files as arrays and tokenize the targets.
   def speech_file_to_array_fn(snake_case__			:       Optional[int]      ):
      UpperCamelCase      ,	UpperCamelCase       :    List[str]		      =      torchaudio.load(batch['path']      )
      UpperCamelCase       :    List[str]		      =      resampler(lowercase_      ).squeeze().numpy()
      UpperCamelCase       :    List[str]		      =      16000
      UpperCamelCase       :    Dict		      =      batch['text']
      return batch

   UpperCamelCase       :    Optional[Any]		      =      train_dataset.map(
       lowercase_			,	remove_columns=train_dataset.column_names			,	num_proc=data_args.preprocessing_num_workers			,	)
   UpperCamelCase       :    Optional[Any]		      =      eval_dataset.map(
       lowercase_			,	remove_columns=eval_dataset.column_names			,	num_proc=data_args.preprocessing_num_workers			,	)

   def prepare_dataset(snake_case__			:       str      ):
      # check that all files have the correct sampling rate
      assert (
          len(set(batch['sampling_rate']      )      ) == 1
      ), F"""Make sure all inputs have the same sampling rate of {processor.feature_extractor.sampling_rate}."""

      UpperCamelCase       :    Optional[Any]		      =      processor(
          audio=batch['speech']			,	text=batch['target_text']			,	sampling_rate=batch['sampling_rate'][0]      )
      batch.update(lowercase_      )
      return batch

   UpperCamelCase       :    int		      =      train_dataset.map(
       lowercase_			,	remove_columns=train_dataset.column_names			,	batch_size=training_args.per_device_train_batch_size			,	batched=lowercase_			,	num_proc=data_args.preprocessing_num_workers			,	)
   UpperCamelCase       :    Dict		      =      eval_dataset.map(
       lowercase_			,	remove_columns=eval_dataset.column_names			,	batch_size=training_args.per_device_train_batch_size			,	batched=lowercase_			,	num_proc=data_args.preprocessing_num_workers			,	)

   # Metric
   UpperCamelCase       :    List[str]		      =      datasets.load_metric('wer'      )

   def compute_metrics(snake_case__			:       Any      ):
      UpperCamelCase       :    Optional[Any]		      =      pred.predictions
      UpperCamelCase       :    Optional[int]		      =      np.argmax(lowercase_			,	axis=-1      )

      UpperCamelCase       :    Tuple		      =      processor.tokenizer.pad_token_id

      UpperCamelCase       :    Optional[Any]		      =      processor.batch_decode(lowercase_      )
      # we do not want to group tokens when computing the metrics
      UpperCamelCase       :    Tuple		      =      processor.batch_decode(pred.label_ids			,	group_tokens=lowercase_      )

      UpperCamelCase       :    Union[str, Any]		      =      wer_metric.compute(predictions=lowercase_			,	references=lowercase_      )

      return {"wer": wer}

   if model_args.freeze_feature_extractor:
      model.freeze_feature_extractor()

   # Data collator
   UpperCamelCase       :    Optional[Any]		      =      DataCollatorCTCWithPadding(processor=lowercase_			,	padding=lowercase_      )

   # Initialize our Trainer
   UpperCamelCase       :    Dict		      =      CTCTrainer(
       model=lowercase_			,	data_collator=lowercase_			,	args=lowercase_			,	compute_metrics=lowercase_			,	train_dataset=train_dataset if training_args.do_train else None			,	eval_dataset=eval_dataset if training_args.do_eval else None			,	tokenizer=processor.feature_extractor			,	)

   # Training
   if training_args.do_train:
      if last_checkpoint is not None:
         UpperCamelCase       :    Dict		      =      last_checkpoint
      elif os.path.isdir(model_args.model_name_or_path      ):
         UpperCamelCase       :    Dict		      =      model_args.model_name_or_path
      else:
         UpperCamelCase       :    Any		      =      None

      # Save the feature_extractor and the tokenizer
      if is_main_process(training_args.local_rank      ):
         processor.save_pretrained(training_args.output_dir      )

      UpperCamelCase       :    Dict		      =      trainer.train(resume_from_checkpoint=lowercase_      )
      trainer.save_model()

      UpperCamelCase       :    int		      =      train_result.metrics
      UpperCamelCase       :    List[str]		      =      (
          data_args.max_train_samples if data_args.max_train_samples is not None else len(lowercase_      )
      )
      UpperCamelCase       :    Any		      =      min(lowercase_			,	len(lowercase_      )      )

      trainer.log_metrics('train'			,	lowercase_      )
      trainer.save_metrics('train'			,	lowercase_      )
      trainer.save_state()

   # Evaluation
   UpperCamelCase       :    List[str]		      =      {}
   if training_args.do_eval:
      logger.info('*** Evaluate ***'      )
      UpperCamelCase       :    Optional[int]		      =      trainer.evaluate()
      UpperCamelCase       :    Optional[int]		      =      data_args.max_val_samples if data_args.max_val_samples is not None else len(lowercase_      )
      UpperCamelCase       :    Dict		      =      min(lowercase_			,	len(lowercase_      )      )

      trainer.log_metrics('eval'			,	lowercase_      )
      trainer.save_metrics('eval'			,	lowercase_      )

   return results


if __name__ == "__main__":
    main()











'''simple docstring'''







import enum
import warnings

from .. import MODEL_FOR_CAUSAL_LM_MAPPING, TF_MODEL_FOR_CAUSAL_LM_MAPPING
from ..utils import add_end_docstrings, is_tf_available
from .base import PIPELINE_INIT_ARGS, Pipeline


if is_tf_available():
  import tensorflow as tf




class   __SCREAMING_SNAKE_CASE       (enum.Enum       ):


  """simple docstring"""


  _a  : Tuple      			=     0
  _a  : List[str]      			=     1
  _a  : int      			=     2





@add_end_docstrings(__A       )
class   __SCREAMING_SNAKE_CASE       (__A       ):


  """simple docstring"""


  _a  : str      			=     '''
    In 1991, the remains of Russian Tsar Nicholas II and his family (except for Alexei and Maria) are discovered. The
    voice of Nicholas\'s young son, Tsarevich Alexei Nikolaevich, narrates the remainder of the story. 1883 Western
    Siberia, a young Grigori Rasputin is asked by his father and a group of men to perform magic. Rasputin has a vision
    and denounces one of the men as a horse thief. Although his father initially slaps him for making such an
    accusation, Rasputin watches as the man is chased outside and beaten. Twenty years later, Rasputin sees a vision of
    the Virgin Mary, prompting him to become a priest. Rasputin quickly becomes famous, with people, even a bishop,
    begging for his blessing. <eod> </s> <eos>
    '''







  def __init__(			self   ,	*UpperCamelCase__   ,	**UpperCamelCase__	):






         """simple docstring"""


         super().__init__(*UpperCamelCase__   ,	**UpperCamelCase__	)
         self.check_model_type(
             TF_MODEL_FOR_CAUSAL_LM_MAPPING if self.framework == 'tf' else MODEL_FOR_CAUSAL_LM_MAPPING	)
         if "prefix" not in self._preprocess_params:
                # This is very specific. The logic is quite complex and needs to be done
                # as a "default".
                # It also defines both some preprocess_kwargs and generate_kwargs
                # which is why we cannot put them in their respective methods.
                a_						     =			None
                if self.model.config.prefix is not None:
                       a_						     =			self.model.config.prefix
                if prefix is None and self.model.__class__.__name__ in [
                    "XLNetLMHeadModel",
                    "TransfoXLLMHeadModel",
                    "TFXLNetLMHeadModel",
                    "TFTransfoXLLMHeadModel",
                ]:
                       # For XLNet and TransformerXL we add an article to the prompt to give more state to the model.
                       a_						     =			self.XL_PREFIX
                if prefix is not None:
                       # Recalculate some generate_kwargs linked to prefix.
                       a_     ,	a_     ,	a_						     =			self._sanitize_parameters(prefix=UpperCamelCase__   ,	**self._forward_params	)
                       a_						     =			{**self._preprocess_params, **preprocess_params}
                       a_						     =			{**self._forward_params, **forward_params}







  def    _a       (			self   ,	UpperCamelCase__=None   ,	UpperCamelCase__=None   ,	UpperCamelCase__=None   ,	UpperCamelCase__=None   ,	UpperCamelCase__=None   ,	UpperCamelCase__=None   ,	UpperCamelCase__=None   ,	UpperCamelCase__=None   ,	**UpperCamelCase__   ,	):






         """simple docstring"""


         a_						     =			{}
         if prefix is not None:
                a_						     =			prefix
         if prefix:
                a_						     =			self.tokenizer(
                    UpperCamelCase__   ,	padding=UpperCamelCase__   ,	add_special_tokens=UpperCamelCase__   ,	return_tensors=self.framework	)
                a_						     =			prefix_inputs['input_ids'].shape[-1]

         if handle_long_generation is not None:
                if handle_long_generation not in {"hole"}:
                       raise ValueError(
                           f'{handle_long_generation} is not a valid value for `handle_long_generation` parameter expected'
                           ' [None, \'hole\']'	)
                a_						     =			handle_long_generation

         preprocess_params.update(UpperCamelCase__	)
         a_						     =			generate_kwargs

         a_						     =			{}
         if return_full_text is not None and return_type is None:
                if return_text is not None:
                       raise ValueError('`return_text` is mutually exclusive with `return_full_text`'	)
                if return_tensors is not None:
                       raise ValueError('`return_full_text` is mutually exclusive with `return_tensors`'	)
                a_						     =			ReturnType.FULL_TEXT if return_full_text else ReturnType.NEW_TEXT
         if return_tensors is not None and return_type is None:
                if return_text is not None:
                       raise ValueError('`return_text` is mutually exclusive with `return_tensors`'	)
                a_						     =			ReturnType.TENSORS
         if return_type is not None:
                a_						     =			return_type
         if clean_up_tokenization_spaces is not None:
                a_						     =			clean_up_tokenization_spaces

         if stop_sequence is not None:
                a_						     =			self.tokenizer.encode(UpperCamelCase__   ,	add_special_tokens=UpperCamelCase__	)
                if len(UpperCamelCase__	) > 1:
                       warnings.warn(
                           'Stopping on a multiple token sequence is not yet supported on transformers. The first token of'
                           ' the stop sequence will be used as the stop sequence string in the interim.'	)
                a_						     =			stop_sequence_ids[0]

         return preprocess_params, forward_params, postprocess_params







  def    _a       (			self   ,	*UpperCamelCase__   ,	**UpperCamelCase__	):






         """simple docstring"""


         if self.model.__class__.__name__ in ["TransfoXLLMHeadModel"]:
                kwargs.update({'add_space_before_punct_symbol': True}	)

         return super()._parse_and_tokenize(*UpperCamelCase__   ,	**UpperCamelCase__	)







  def __call__(			self   ,	UpperCamelCase__   ,	**UpperCamelCase__	):






         """simple docstring"""


         return super().__call__(UpperCamelCase__   ,	**UpperCamelCase__	)







  def    _a       (			self   ,	UpperCamelCase__   ,	UpperCamelCase__=""   ,	UpperCamelCase__=None   ,	**UpperCamelCase__	):






         """simple docstring"""


         a_						     =			self.tokenizer(
             prefix + prompt_text   ,	padding=UpperCamelCase__   ,	add_special_tokens=UpperCamelCase__   ,	return_tensors=self.framework	)
         a_						     =			prompt_text

         if handle_long_generation == "hole":
                a_						     =			inputs['input_ids'].shape[-1]
                if "max_new_tokens" in generate_kwargs:
                       a_						     =			generate_kwargs['max_new_tokens']
                else:
                       a_						     =			generate_kwargs.get('max_length'   ,	self.model.config.max_length	) - cur_len
                       if new_tokens < 0:
                              raise ValueError('We cannot infer how many new tokens are expected'	)
                if cur_len + new_tokens > self.tokenizer.model_max_length:
                       a_						     =			self.tokenizer.model_max_length - new_tokens
                       if keep_length <= 0:
                              raise ValueError(
                                  'We cannot use `hole` to handle this generation the number of desired tokens exceeds the'
                                  ' models max length'	)

                       a_						     =			inputs['input_ids'][:, -keep_length:]
                       if "attention_mask" in inputs:
                              a_						     =			inputs['attention_mask'][:, -keep_length:]

         return inputs







  def    _a       (			self   ,	UpperCamelCase__   ,	**UpperCamelCase__	):






         """simple docstring"""


         a_						     =			model_inputs['input_ids']
         a_						     =			model_inputs.get('attention_mask'   ,	UpperCamelCase__	)
         # Allow empty prompts
         if input_ids.shape[1] == 0:
                a_						     =			None
                a_						     =			None
                a_						     =			1
         else:
                a_						     =			input_ids.shape[0]
         a_						     =			model_inputs.pop('prompt_text'	)

         # If there is a prefix, we may need to adjust the generation length. Do so without permanently modifying
         # generate_kwargs, as some of the parameterization may come from the initialization of the pipeline.
         a_						     =			generate_kwargs.pop('prefix_length'   ,	0	)
         if prefix_length > 0:
                a_						     =			'max_new_tokens' in generate_kwargs or (
                    'generation_config' in generate_kwargs
                    and generate_kwargs['generation_config'].max_new_tokens is not None
                )
                if not has_max_new_tokens:
                       a_						     =			generate_kwargs.get('max_length'	) or self.model.config.max_length
                       generate_kwargs["max_length"] += prefix_length
                a_						     =			'min_new_tokens' in generate_kwargs or (
                    'generation_config' in generate_kwargs
                    and generate_kwargs['generation_config'].min_new_tokens is not None
                )
                if not has_min_new_tokens and "min_length" in generate_kwargs:
                       generate_kwargs["min_length"] += prefix_length

        # BS x SL
         a_						     =			self.model.generate(input_ids=UpperCamelCase__   ,	attention_mask=UpperCamelCase__   ,	**UpperCamelCase__	)
         a_						     =			generated_sequence.shape[0]
         if self.framework == "pt":
                a_						     =			generated_sequence.reshape(UpperCamelCase__   ,	out_b // in_b   ,	*generated_sequence.shape[1:]	)
         elif self.framework == "tf":
                a_						     =			tf.reshape(UpperCamelCase__   ,	(in_b, out_b // in_b, *generated_sequence.shape[1:])	)
         return {"generated_sequence": generated_sequence, "input_ids": input_ids, "prompt_text": prompt_text}







  def    _a       (			self   ,	UpperCamelCase__   ,	UpperCamelCase__=ReturnType.FULL_TEXT   ,	UpperCamelCase__=True	):






         """simple docstring"""


         a_						     =			model_outputs['generated_sequence'][0]
         a_						     =			model_outputs['input_ids']
         a_						     =			model_outputs['prompt_text']
         a_						     =			generated_sequence.numpy().tolist()
         a_						     =			[]
         for sequence in generated_sequence:
                if return_type == ReturnType.TENSORS:
                       a_						     =			{'generated_token_ids': sequence}
                elif return_type in {ReturnType.NEW_TEXT, ReturnType.FULL_TEXT}:
                       # Decode text
                       a_						     =			self.tokenizer.decode(
                           UpperCamelCase__   ,	skip_special_tokens=UpperCamelCase__   ,	clean_up_tokenization_spaces=UpperCamelCase__   ,	)

                       # Remove PADDING prompt of the sequence if XLNet or Transfo-XL model is used
                       if input_ids is None:
                              a_						     =			0
                       else:
                              a_						     =			len(
                                  self.tokenizer.decode(
                                      input_ids[0]   ,	skip_special_tokens=UpperCamelCase__   ,	clean_up_tokenization_spaces=UpperCamelCase__   ,	)	)

                       if return_type == ReturnType.FULL_TEXT:
                              a_						     =			prompt_text + text[prompt_length:]
                       else:
                              a_						     =			text[prompt_length:]

                       a_						     =			{'generated_text': all_text}
                records.append(UpperCamelCase__	)

         return records




