




import logging
import pathlib
import re
import sys
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, List, Optional, Set, Union

import datasets
import librosa
import numpy as np
import torch
from lang_trans import arabic
from packaging import version
from torch import nn

from transformers import (
    HfArgumentParser,
    Trainer,
    TrainingArguments,
    WavaVecaCTCTokenizer,
    WavaVecaFeatureExtractor,
    WavaVecaForCTC,
    WavaVecaProcessor,
    is_apex_available,
    trainer_utils,
)


if is_apex_available():
       from apex import amp

if version.parse(version.parse(torch.__version__).base_version) >= version.parse('1.6'):
       SCREAMING_SNAKE_CASE__ :   Union[str, Any]						=     True
       from torch.cuda.amp import autocast


SCREAMING_SNAKE_CASE__ :   Any						=     logging.getLogger(__name__)




@dataclass
class 	UpperCamelCase__			:

    '''simple docstring'''



    lowerCamelCase_   :			str			  =   field(
        metadata={"""help""": """Path to pretrained model or model identifier from huggingface.co/models"""}  )
    lowerCamelCase_   :			Optional[str]			  =   field(
        default=lowerCAmelCase__					,       metadata={"""help""": """Where do you want to store the pretrained models downloaded from huggingface.co"""}					,       )
    lowerCamelCase_   :			Optional[bool]			  =   field(
        default=lowerCAmelCase__					,       metadata={"""help""": """Whether to freeze the feature extractor layers of the model."""}  )
    lowerCamelCase_   :			Optional[bool]			  =   field(
        default=lowerCAmelCase__					,       metadata={"""help""": """Whether to log verbose messages or not."""}					,       )




def   A   (   _SCREAMING_SNAKE_CASE				,_SCREAMING_SNAKE_CASE					)			->				Dict:
       logging.basicConfig(
           format="%(asctime)s - %(levelname)s - %(name)s - %(message)s"				,datefmt="%m/%d/%Y %H:%M:%S"				,handlers=[logging.StreamHandler(sys.stdout					)]				,)
       lowerCamelCase     :   str    				=      logging.WARNING
       if model_args.verbose_logging:
              lowerCamelCase     :   int    				=      logging.DEBUG
       elif trainer_utils.is_main_process(training_args.local_rank					):
              lowerCamelCase     :   Union[str, Any]    				=      logging.INFO
       logger.setLevel(_SCREAMING_SNAKE_CASE					)




@dataclass
class 	UpperCamelCase__			:

    '''simple docstring'''



    lowerCamelCase_   :			str			  =   field(
        default=lowerCAmelCase__					,       metadata={"""help""": """The name of the dataset to use (via the datasets library)."""}  )
    lowerCamelCase_   :			Optional[str]			  =   field(
        default=lowerCAmelCase__					,       metadata={"""help""": """The configuration name of the dataset to use (via the datasets library)."""}  )
    lowerCamelCase_   :			Optional[str]			  =   field(
        default="""train"""					,       metadata={
            """help""": """The name of the training data set split to use (via the datasets library). Defaults to 'train'"""
        }					,       )
    lowerCamelCase_   :			Optional[str]			  =   field(
        default="""validation"""					,       metadata={
            """help""": (
                """The name of the validation data set split to use (via the datasets library). Defaults to 'validation'"""
            )
        }					,       )
    lowerCamelCase_   :			Optional[str]			  =   field(
        default="""text"""					,       metadata={"""help""": """Column in the dataset that contains label (target text). Defaults to 'text'"""}					,       )
    lowerCamelCase_   :			Optional[str]			  =   field(
        default="""file"""					,       metadata={"""help""": """Column in the dataset that contains speech file path. Defaults to 'file'"""}					,       )
    lowerCamelCase_   :			Optional[bool]			  =   field(
        default=lowerCAmelCase__					,       metadata={"""help""": """Resample loaded audio to target feature extractor's sampling rate or not."""}					,       )
    lowerCamelCase_   :			Optional[float]			  =   field(
        default=lowerCAmelCase__					,       metadata={"""help""": """Filters out examples longer than specified. Defaults to no filtering."""}					,       )
    lowerCamelCase_   :			Optional[str]			  =   field(
        default="""librispeech"""					,       metadata={
            """help""": (
                """Orthography used for normalization and tokenization: 'librispeech' (default), 'timit', or"""
                """ 'buckwalter'."""
            )
        }					,       )
    lowerCamelCase_   :			bool			  =   field(
        default=lowerCAmelCase__					,       metadata={"""help""": """Overwrite the cached preprocessed datasets or not."""}  )
    lowerCamelCase_   :			Optional[int]			  =   field(
        default=lowerCAmelCase__					,       metadata={"""help""": """The number of processes to use for the preprocessing."""}					,       )




@dataclass
class 	UpperCamelCase__			:

    '''simple docstring'''



    lowerCamelCase_   :			bool			  =   False
    lowerCamelCase_   :			Optional[str]			  =   None
    lowerCamelCase_   :			Optional[str]			  =   "|"
    lowerCamelCase_   :			Optional[Dict[str, str]]			  =   field(default_factory=lowerCAmelCase__  )
    lowerCamelCase_   :			Optional[Set[str]]			  =   field(default_factory=lowerCAmelCase__  )
    lowerCamelCase_   :			Optional[Callable[[str], str]]			  =   None







    @classmethod
    def 					_lowercase  (					cls  ,						UpperCamelCase__  )  ->   str:
           if name == "librispeech":
                  return cls()
           if name == "timit":
                  return cls(
                      do_lower_case=UpperCamelCase__  ,						translation_table=str.maketrans({"-": " "}  )  ,						)
           if name == "buckwalter":
                  lowerCamelCase     :   List[str]    				=      {
                      "-": " ",  # sometimes used to represent pauses
                      "^": "v",  # fixing "tha" in arabic_speech_corpus dataset
                  }
                  return cls(
                      vocab_file=pathlib.Path(UpperCamelCase__  ).parent.joinpath("vocab/buckwalter.json"  )  ,						word_delimiter_token="/"  ,						translation_table=str.maketrans(UpperCamelCase__  )  ,						words_to_remove={"sil"}  ,						untransliterator=arabic.buckwalter.untransliterate  ,						)
           raise ValueError(F'''Unsupported orthography: \'{name}\'.'''  )







    def 					_lowercase  (					self  ,						UpperCamelCase__  )  ->   str:
           # TODO(elgeish) return a pipeline (e.g., from jiwer) instead? Or rely on branch predictor as is
           if len(self.translation_table  ) > 0:
                  lowerCamelCase     :   Any    				=      text.translate(self.translation_table  )
           if len(self.words_to_remove  ) == 0:
                  lowerCamelCase     :   List[Any]    				=      " ".join(text.split()  )  # clean up whitespaces
           else:
                  lowerCamelCase     :   List[Any]    				=      " ".join(w for w in text.split() if w not in self.words_to_remove  )  # and clean up whilespaces
           return text







    def 					_lowercase  (					self  ,						UpperCamelCase__  )  ->   WavaVecaProcessor:
           lowerCamelCase     :   Union[str, Any]    				=      WavaVecaFeatureExtractor.from_pretrained(
               model_args.model_name_or_path  ,						cache_dir=model_args.cache_dir  )
           if self.vocab_file:
                  lowerCamelCase     :   str    				=      WavaVecaCTCTokenizer(
                      self.vocab_file  ,						cache_dir=model_args.cache_dir  ,						do_lower_case=self.do_lower_case  ,						word_delimiter_token=self.word_delimiter_token  ,						)
           else:
                  lowerCamelCase     :   List[str]    				=      WavaVecaCTCTokenizer.from_pretrained(
                      model_args.model_name_or_path  ,						cache_dir=model_args.cache_dir  ,						do_lower_case=self.do_lower_case  ,						word_delimiter_token=self.word_delimiter_token  ,						)
           return WavaVecaProcessor(UpperCamelCase__  ,						UpperCamelCase__  )




@dataclass
class 	UpperCamelCase__			:

    '''simple docstring'''



    lowerCamelCase_   :			WavaVecaProcessor
    lowerCamelCase_   :			Union[bool, str]			  =   True
    lowerCamelCase_   :			Optional[int]			  =   None
    lowerCamelCase_   :			Optional[int]			  =   None
    lowerCamelCase_   :			Optional[int]			  =   None
    lowerCamelCase_   :			Optional[int]			  =   None







    def __call__(					self  ,						UpperCamelCase__  )  ->   Dict[str, torch.Tensor]:
           # split inputs and labels since they have to be of different lenghts and need
           # different padding methods
           lowerCamelCase     :   Dict    				=      [{"input_values": feature["input_values"]} for feature in features]
           lowerCamelCase     :   Tuple    				=      [{"input_ids": feature["labels"]} for feature in features]

           lowerCamelCase     :   str    				=      self.processor.pad(
               UpperCamelCase__  ,						padding=self.padding  ,						max_length=self.max_length  ,						pad_to_multiple_of=self.pad_to_multiple_of  ,						return_tensors="pt"  ,						)
           lowerCamelCase     :   Tuple    				=      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
           lowerCamelCase     :   int    				=      labels_batch["input_ids"].masked_fill(labels_batch.attention_mask.ne(1  )  ,						-100  )

           lowerCamelCase     :   str    				=      labels

           return batch




class 	UpperCamelCase__			(lowerCAmelCase__  ):

    '''simple docstring'''



    def 					_lowercase  (					self  ,						UpperCamelCase__  ,						UpperCamelCase__  )  ->   torch.Tensor:
           model.train()
           lowerCamelCase     :   Tuple    				=      self._prepare_inputs(UpperCamelCase__  )

           if self.use_amp:
                  with autocast():
                         lowerCamelCase     :   Optional[int]    				=      self.compute_loss(UpperCamelCase__  ,						UpperCamelCase__  )
           else:
                  lowerCamelCase     :   int    				=      self.compute_loss(UpperCamelCase__  ,						UpperCamelCase__  )

           if self.args.n_gpu > 1:
                  if model.module.config.ctc_loss_reduction == "mean":
                         lowerCamelCase     :   List[str]    				=      loss.mean()
                  elif model.module.config.ctc_loss_reduction == "sum":
                         lowerCamelCase     :   List[str]    				=      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:
                  lowerCamelCase     :   Tuple    				=      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   A   (   )			->				List[Any]:
       # See all possible arguments in src/transformers/training_args.py
       # or by passing the --help flag to this script.
       # We now keep distinct sets of args, for a cleaner separation of concerns.

       lowerCamelCase     :   str    				=      HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)					)

       lowerCamelCase				,   lowerCamelCase				,   lowerCamelCase     :   Union[str, Any]    				=      parser.parse_args_into_dataclasses()
       configure_logger(_SCREAMING_SNAKE_CASE				,_SCREAMING_SNAKE_CASE					)

       lowerCamelCase     :   Optional[int]    				=      Orthography.from_name(data_args.orthography.lower()					)
       lowerCamelCase     :   int    				=      orthography.create_processor(_SCREAMING_SNAKE_CASE					)
       lowerCamelCase     :   Optional[int]    				=      WavaVecaForCTC.from_pretrained(
           model_args.model_name_or_path				,cache_dir=model_args.cache_dir				,gradient_checkpointing=training_args.gradient_checkpointing				,vocab_size=len(processor.tokenizer					)				,)

       lowerCamelCase     :   Tuple    				=      datasets.load_dataset(
           data_args.dataset_name				,data_args.dataset_config_name				,split=data_args.train_split_name					)
       lowerCamelCase     :   List[Any]    				=      datasets.load_dataset(
           data_args.dataset_name				,data_args.dataset_config_name				,split=data_args.validation_split_name					)

       lowerCamelCase     :   List[str]    				=      datasets.load_metric("wer"					)
       lowerCamelCase     :   List[str]    				=      processor.feature_extractor.sampling_rate if data_args.target_feature_extractor_sampling_rate else None
       lowerCamelCase     :   Union[str, Any]    				=      "".join(t for t in processor.tokenizer.get_vocab().keys() if len(_SCREAMING_SNAKE_CASE					) == 1					)
       lowerCamelCase     :   Tuple    				=      re.compile(  # remove characters not in vocabulary
           rf'''[^\s{re.escape(_SCREAMING_SNAKE_CASE					)}]'''				,flags=re.IGNORECASE if processor.tokenizer.do_lower_case else 0				,)
       lowerCamelCase     :   str    				=      []

       def prepare_example(_SCREAMING_SNAKE_CASE					):  # TODO(elgeish) make use of multiprocessing?
              lowerCamelCase				,   lowerCamelCase     :   Optional[int]    				=      librosa.load(example[data_args.speech_file_column]				,sr=_SCREAMING_SNAKE_CASE					)
              if data_args.max_duration_in_seconds is not None:
                     lowerCamelCase     :   int    				=      len(example["speech"]					) / example["sampling_rate"]
              # Normalize and clean up text; order matters!
              lowerCamelCase     :   int    				=      orthography.preprocess_for_training(example[data_args.target_text_column]					)
              lowerCamelCase     :   Union[str, Any]    				=      vocabulary_text_cleaner.sub(""				,_SCREAMING_SNAKE_CASE					)
              if updated_text != example[data_args.target_text_column]:
                     text_updates.append((example[data_args.target_text_column], updated_text)					)
                     lowerCamelCase     :   Dict    				=      updated_text
              return example

       lowerCamelCase     :   Optional[int]    				=      train_dataset.map(_SCREAMING_SNAKE_CASE				,remove_columns=[data_args.speech_file_column]					)
       lowerCamelCase     :   Tuple    				=      val_dataset.map(_SCREAMING_SNAKE_CASE				,remove_columns=[data_args.speech_file_column]					)

       if data_args.max_duration_in_seconds is not None:

              def filter_by_max_duration(_SCREAMING_SNAKE_CASE					):
                     return example["duration_in_seconds"] <= data_args.max_duration_in_seconds

              lowerCamelCase     :   Optional[int]    				=      len(_SCREAMING_SNAKE_CASE					)
              lowerCamelCase     :   Dict    				=      len(_SCREAMING_SNAKE_CASE					)
              lowerCamelCase     :   Dict    				=      train_dataset.filter(_SCREAMING_SNAKE_CASE				,remove_columns=["duration_in_seconds"]					)
              lowerCamelCase     :   Dict    				=      val_dataset.filter(_SCREAMING_SNAKE_CASE				,remove_columns=["duration_in_seconds"]					)
              if len(_SCREAMING_SNAKE_CASE					) > old_train_size:
                     logger.warning(
                         f'''Filtered out {len(_SCREAMING_SNAKE_CASE					) - old_train_size} train example(s) longer than'''
                         f''' {data_args.max_duration_in_seconds} second(s).'''					)
              if len(_SCREAMING_SNAKE_CASE					) > old_val_size:
                     logger.warning(
                         f'''Filtered out {len(_SCREAMING_SNAKE_CASE					) - old_val_size} validation example(s) longer than'''
                         f''' {data_args.max_duration_in_seconds} second(s).'''					)
       logger.info(f'''Split sizes: {len(_SCREAMING_SNAKE_CASE					)} train and {len(_SCREAMING_SNAKE_CASE					)} validation.'''					)

       logger.warning(f'''Updated {len(_SCREAMING_SNAKE_CASE					)} transcript(s) using \'{data_args.orthography}\' orthography rules.'''					)
       if logger.isEnabledFor(logging.DEBUG					):
              for original_text, updated_text in text_updates:
                     logger.debug(f'''Updated text: "{original_text}" -> "{updated_text}"'''					)
       lowerCamelCase     :   Optional[int]    				=      None

       def prepare_dataset(_SCREAMING_SNAKE_CASE					):
              # 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}.'''

              lowerCamelCase     :   Optional[int]    				=      processor(
                  audio=batch["speech"]				,text=batch[data_args.target_text_column]				,sampling_rate=batch["sampling_rate"][0]					)
              batch.update(_SCREAMING_SNAKE_CASE					)
              return batch

       lowerCamelCase     :   Optional[Any]    				=      train_dataset.map(
           _SCREAMING_SNAKE_CASE				,batch_size=training_args.per_device_train_batch_size				,batched=_SCREAMING_SNAKE_CASE				,num_proc=data_args.preprocessing_num_workers				,)
       lowerCamelCase     :   Optional[Any]    				=      val_dataset.map(
           _SCREAMING_SNAKE_CASE				,batch_size=training_args.per_device_train_batch_size				,batched=_SCREAMING_SNAKE_CASE				,num_proc=data_args.preprocessing_num_workers				,)

       lowerCamelCase     :   Optional[Any]    				=      DataCollatorCTCWithPadding(processor=_SCREAMING_SNAKE_CASE				,padding=_SCREAMING_SNAKE_CASE					)

       def compute_metrics(_SCREAMING_SNAKE_CASE					):
              lowerCamelCase     :   Dict    				=      pred.predictions
              lowerCamelCase     :   int    				=      np.argmax(_SCREAMING_SNAKE_CASE				,axis=-1					)

              lowerCamelCase     :   Any    				=      processor.tokenizer.pad_token_id

              lowerCamelCase     :   Tuple    				=      processor.batch_decode(_SCREAMING_SNAKE_CASE					)
              # we do not want to group tokens when computing the metrics
              lowerCamelCase     :   Optional[int]    				=      processor.batch_decode(pred.label_ids				,group_tokens=_SCREAMING_SNAKE_CASE					)
              if logger.isEnabledFor(logging.DEBUG					):
                     for reference, predicted in zip(_SCREAMING_SNAKE_CASE				,_SCREAMING_SNAKE_CASE					):
                            logger.debug(f'''reference: "{reference}"'''					)
                            logger.debug(f'''predicted: "{predicted}"'''					)
                            if orthography.untransliterator is not None:
                                   logger.debug(f'''reference (untransliterated): "{orthography.untransliterator(_SCREAMING_SNAKE_CASE					)}"'''					)
                                   logger.debug(f'''predicted (untransliterated): "{orthography.untransliterator(_SCREAMING_SNAKE_CASE					)}"'''					)

              lowerCamelCase     :   Any    				=      wer_metric.compute(predictions=_SCREAMING_SNAKE_CASE				,references=_SCREAMING_SNAKE_CASE					)

              return {"wer": wer}

       if model_args.freeze_feature_extractor:
              model.freeze_feature_extractor()

       lowerCamelCase     :   int    				=      CTCTrainer(
           model=_SCREAMING_SNAKE_CASE				,data_collator=_SCREAMING_SNAKE_CASE				,args=_SCREAMING_SNAKE_CASE				,compute_metrics=_SCREAMING_SNAKE_CASE				,train_dataset=_SCREAMING_SNAKE_CASE				,eval_dataset=_SCREAMING_SNAKE_CASE				,tokenizer=processor.feature_extractor				,)

       trainer.train()


if __name__ == "__main__":
       main()












import collections
import types

import numpy as np

from ..utils import (
    add_end_docstrings,
    is_tensorflow_probability_available,
    is_tf_available,
    is_torch_available,
    requires_backends,
)
from .base import PIPELINE_INIT_ARGS, ArgumentHandler, Dataset, Pipeline, PipelineException


if is_torch_available():
       import torch

       from ..models.auto.modeling_auto import (
           MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,
           MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING,
       )

if is_tf_available() and is_tensorflow_probability_available():
       import tensorflow as tf
       import tensorflow_probability as tfp

       from ..models.auto.modeling_tf_auto import (
           TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,
           TF_MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING,
       )




class 	UpperCamelCase__			(lowerCAmelCase__  ):

    '''simple docstring'''



    def __call__(					self  ,						UpperCamelCase__=None  ,						UpperCamelCase__=None  ,						**UpperCamelCase__  )  ->   List[Any]:
           # Returns tqa_pipeline_inputs of shape:
           # [
           #   {"table": pd.DataFrame, "query": List[str]},
           #   ...,
           #   {"table": pd.DataFrame, "query" : List[str]}
           # ]
           requires_backends(self  ,						"pandas"  )
           import pandas as pd

           if table is None:
                  raise ValueError("Keyword argument `table` cannot be None."  )
           elif query is None:
                  if isinstance(UpperCamelCase__  ,						UpperCamelCase__  ) and table.get("query"  ) is not None and table.get("table"  ) is not None:
                         lowerCamelCase     :   Optional[int]    				=      [table]
                  elif isinstance(UpperCamelCase__  ,						UpperCamelCase__  ) and len(UpperCamelCase__  ) > 0:
                         if not all(isinstance(UpperCamelCase__  ,						UpperCamelCase__  ) for d in table  ):
                                raise ValueError(
                                    F'''Keyword argument `table` should be a list of dict, but is {(type(UpperCamelCase__  ) for d in table)}'''  )

                         if table[0].get("query"  ) is not None and table[0].get("table"  ) is not None:
                                lowerCamelCase     :   int    				=      table
                         else:
                                raise ValueError(
                                    "If keyword argument `table` is a list of dictionaries, each dictionary should have a `table`"
                                    F''' and `query` key, but only dictionary has keys {table[0].keys()} `table` and `query` keys.'''  )
                  elif Dataset is not None and isinstance(UpperCamelCase__  ,						UpperCamelCase__  ) or isinstance(UpperCamelCase__  ,						types.GeneratorType  ):
                         return table
                  else:
                         raise ValueError(
                             "Invalid input. Keyword argument `table` should be either of type `dict` or `list`, but "
                             F'''is {type(UpperCamelCase__  )})'''  )
           else:
                  lowerCamelCase     :   Tuple    				=      [{"table": table, "query": query}]

           for tqa_pipeline_input in tqa_pipeline_inputs:
                  if not isinstance(tqa_pipeline_input["table"]  ,						pd.DataFrame  ):
                         if tqa_pipeline_input["table"] is None:
                                raise ValueError("Table cannot be None."  )

                         lowerCamelCase     :   Optional[Any]    				=      pd.DataFrame(tqa_pipeline_input["table"]  )

           return tqa_pipeline_inputs




@add_end_docstrings(lowerCAmelCase__  )
class 	UpperCamelCase__			(lowerCAmelCase__  ):

    '''simple docstring'''



    lowerCamelCase_   :			Optional[Any]			  =   """table,query"""







    def __init__(					self  ,						UpperCamelCase__=TableQuestionAnsweringArgumentHandler()  ,						*UpperCamelCase__  ,						**UpperCamelCase__  )  ->   Dict:
           super().__init__(*UpperCamelCase__  ,						**UpperCamelCase__  )
           lowerCamelCase     :   Any    				=      args_parser

           self.check_model_type(
               dict(
                   TF_MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING.items()
                   + TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING.items()  )
               if self.framework == "tf"
               else dict(
                   MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING.items() + MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING.items()  )  )

           lowerCamelCase     :   Any    				=      bool(getattr(self.model.config  ,						"aggregation_labels"  ,						UpperCamelCase__  )  ) and bool(
               getattr(self.model.config  ,						"num_aggregation_labels"  ,						UpperCamelCase__  )  )
           lowerCamelCase     :   List[Any]    				=      "tapas" if hasattr(self.model.config  ,						"aggregation_labels"  ) else None







    def 					_lowercase  (					self  ,						**UpperCamelCase__  )  ->   int:
           return self.model(**UpperCamelCase__  )







    def 					_lowercase  (					self  ,						**UpperCamelCase__  )  ->   Union[str, Any]:
           if self.framework == "pt":
                  lowerCamelCase     :   List[str]    				=      []
                  lowerCamelCase     :   List[Any]    				=      []
                  lowerCamelCase     :   List[Any]    				=      None
                  lowerCamelCase     :   List[str]    				=      inputs["input_ids"].shape[0]

                  lowerCamelCase     :   Optional[int]    				=      inputs["input_ids"].to(self.device  )
                  lowerCamelCase     :   List[str]    				=      inputs["attention_mask"].to(self.device  )
                  lowerCamelCase     :   str    				=      inputs["token_type_ids"].to(self.device  )
                  lowerCamelCase     :   str    				=      None

                  for index in range(UpperCamelCase__  ):
                         # If sequences have already been processed, the token type IDs will be created according to the previous
                         # answer.
                         if prev_answers is not None:
                                lowerCamelCase     :   List[Any]    				=      token_type_ids_example[:, 3]  # shape (seq_len,)
                                lowerCamelCase     :   Optional[int]    				=      np.zeros_like(prev_labels_example.cpu().numpy()  )  # shape (seq_len,)

                                lowerCamelCase     :   Optional[Any]    				=      token_type_ids[index]  # shape (seq_len, 7)
                                for i in range(model_labels.shape[0]  ):
                                       lowerCamelCase     :   Tuple    				=      token_type_ids_example[:, 0].tolist()[i]
                                       lowerCamelCase     :   Union[str, Any]    				=      token_type_ids_example[:, 1].tolist()[i] - 1
                                       lowerCamelCase     :   Optional[int]    				=      token_type_ids_example[:, 2].tolist()[i] - 1

                                       if row_id >= 0 and col_id >= 0 and segment_id == 1:
                                              lowerCamelCase     :   Dict    				=      int(prev_answers[(col_id, row_id)]  )

                                lowerCamelCase     :   Dict    				=      torch.from_numpy(UpperCamelCase__  ).type(torch.long  ).to(self.device  )

                         lowerCamelCase     :   int    				=      input_ids[index]
                         lowerCamelCase     :   Union[str, Any]    				=      attention_mask[index]  # shape (seq_len,)
                         lowerCamelCase     :   Dict    				=      token_type_ids[index]  # shape (seq_len, 7)
                         lowerCamelCase     :   Dict    				=      self.model(
                             input_ids=input_ids_example.unsqueeze(0  )  ,						attention_mask=attention_mask_example.unsqueeze(0  )  ,						token_type_ids=token_type_ids_example.unsqueeze(0  )  ,						)
                         lowerCamelCase     :   str    				=      outputs.logits

                         if self.aggregate:
                                all_aggregations.append(outputs.logits_aggregation  )

                         all_logits.append(UpperCamelCase__  )

                         lowerCamelCase     :   int    				=      torch.distributions.Bernoulli(logits=UpperCamelCase__  )
                         lowerCamelCase     :   List[Any]    				=      dist_per_token.probs * attention_mask_example.type(torch.floataa  ).to(
                             dist_per_token.probs.device  )

                         lowerCamelCase     :   Optional[int]    				=      collections.defaultdict(UpperCamelCase__  )
                         for i, p in enumerate(probabilities.squeeze().tolist()  ):
                                lowerCamelCase     :   Optional[int]    				=      token_type_ids_example[:, 0].tolist()[i]
                                lowerCamelCase     :   List[Any]    				=      token_type_ids_example[:, 1].tolist()[i] - 1
                                lowerCamelCase     :   str    				=      token_type_ids_example[:, 2].tolist()[i] - 1
                                if col >= 0 and row >= 0 and segment_id == 1:
                                       coords_to_probs[(col, row)].append(UpperCamelCase__  )

                         lowerCamelCase     :   Optional[int]    				=      {key: np.array(coords_to_probs[key]  ).mean() > 0.5 for key in coords_to_probs}

                  lowerCamelCase     :   Optional[Any]    				=      torch.cat(tuple(UpperCamelCase__  )  ,						0  )

                  return (logits_batch,) if not self.aggregate else (logits_batch, torch.cat(tuple(UpperCamelCase__  )  ,						0  ))
           else:
                  lowerCamelCase     :   Optional[Any]    				=      []
                  lowerCamelCase     :   List[Any]    				=      []
                  lowerCamelCase     :   Union[str, Any]    				=      None
                  lowerCamelCase     :   int    				=      inputs["input_ids"].shape[0]

                  lowerCamelCase     :   Dict    				=      inputs["input_ids"]
                  lowerCamelCase     :   Tuple    				=      inputs["attention_mask"]
                  lowerCamelCase     :   Optional[int]    				=      inputs["token_type_ids"].numpy()
                  lowerCamelCase     :   Union[str, Any]    				=      None

                  for index in range(UpperCamelCase__  ):
                         # If sequences have already been processed, the token type IDs will be created according to the previous
                         # answer.
                         if prev_answers is not None:
                                lowerCamelCase     :   Tuple    				=      token_type_ids_example[:, 3]  # shape (seq_len,)
                                lowerCamelCase     :   Optional[Any]    				=      np.zeros_like(UpperCamelCase__  ,						dtype=np.intaa  )  # shape (seq_len,)

                                lowerCamelCase     :   Optional[Any]    				=      token_type_ids[index]  # shape (seq_len, 7)
                                for i in range(model_labels.shape[0]  ):
                                       lowerCamelCase     :   Any    				=      token_type_ids_example[:, 0].tolist()[i]
                                       lowerCamelCase     :   int    				=      token_type_ids_example[:, 1].tolist()[i] - 1
                                       lowerCamelCase     :   Union[str, Any]    				=      token_type_ids_example[:, 2].tolist()[i] - 1

                                       if row_id >= 0 and col_id >= 0 and segment_id == 1:
                                              lowerCamelCase     :   Tuple    				=      int(prev_answers[(col_id, row_id)]  )

                                lowerCamelCase     :   Any    				=      model_labels

                         lowerCamelCase     :   Optional[int]    				=      input_ids[index]
                         lowerCamelCase     :   Optional[Any]    				=      attention_mask[index]  # shape (seq_len,)
                         lowerCamelCase     :   Optional[int]    				=      token_type_ids[index]  # shape (seq_len, 7)
                         lowerCamelCase     :   Union[str, Any]    				=      self.model(
                             input_ids=np.expand_dims(UpperCamelCase__  ,						axis=0  )  ,						attention_mask=np.expand_dims(UpperCamelCase__  ,						axis=0  )  ,						token_type_ids=np.expand_dims(UpperCamelCase__  ,						axis=0  )  ,						)
                         lowerCamelCase     :   Dict    				=      outputs.logits

                         if self.aggregate:
                                all_aggregations.append(outputs.logits_aggregation  )

                         all_logits.append(UpperCamelCase__  )

                         lowerCamelCase     :   Optional[int]    				=      tfp.distributions.Bernoulli(logits=UpperCamelCase__  )
                         lowerCamelCase     :   Any    				=      dist_per_token.probs_parameter() * tf.cast(UpperCamelCase__  ,						tf.floataa  )

                         lowerCamelCase     :   List[Any]    				=      collections.defaultdict(UpperCamelCase__  )
                         lowerCamelCase     :   Optional[int]    				=      token_type_ids_example
                         for i, p in enumerate(tf.squeeze(UpperCamelCase__  ).numpy().tolist()  ):
                                lowerCamelCase     :   Union[str, Any]    				=      token_type_ids_example[:, 0].tolist()[i]
                                lowerCamelCase     :   Any    				=      token_type_ids_example[:, 1].tolist()[i] - 1
                                lowerCamelCase     :   List[str]    				=      token_type_ids_example[:, 2].tolist()[i] - 1
                                if col >= 0 and row >= 0 and segment_id == 1:
                                       coords_to_probs[(col, row)].append(UpperCamelCase__  )

                         lowerCamelCase     :   Optional[Any]    				=      {key: np.array(coords_to_probs[key]  ).mean() > 0.5 for key in coords_to_probs}

                  lowerCamelCase     :   int    				=      tf.concat(tuple(UpperCamelCase__  )  ,						0  )

                  return (logits_batch,) if not self.aggregate else (logits_batch, tf.concat(tuple(UpperCamelCase__  )  ,						0  ))







    def __call__(					self  ,						*UpperCamelCase__  ,						**UpperCamelCase__  )  ->   int:
           lowerCamelCase     :   Union[str, Any]    				=      self._args_parser(*UpperCamelCase__  ,						**UpperCamelCase__  )

           lowerCamelCase     :   Optional[int]    				=      super().__call__(UpperCamelCase__  ,						**UpperCamelCase__  )
           if len(UpperCamelCase__  ) == 1:
                  return results[0]
           return results







    def 					_lowercase  (					self  ,						UpperCamelCase__=None  ,						UpperCamelCase__=None  ,						UpperCamelCase__=None  ,						**UpperCamelCase__  )  ->   Optional[int]:
           lowerCamelCase     :   int    				=      {}
           if padding is not None:
                  lowerCamelCase     :   Dict    				=      padding
           if truncation is not None:
                  lowerCamelCase     :   Tuple    				=      truncation

           lowerCamelCase     :   List[str]    				=      {}
           if sequential is not None:
                  lowerCamelCase     :   Any    				=      sequential
           return preprocess_params, forward_params, {}







    def 					_lowercase  (					self  ,						UpperCamelCase__  ,						UpperCamelCase__=None  ,						UpperCamelCase__=True  ,						UpperCamelCase__=None  )  ->   str:
           if truncation is None:
                  if self.type == "tapas":
                         lowerCamelCase     :   Optional[Any]    				=      "drop_rows_to_fit"
                  else:
                         lowerCamelCase     :   str    				=      "do_not_truncate"

           lowerCamelCase				,   lowerCamelCase     :   Tuple    				=      pipeline_input["table"], pipeline_input["query"]
           if table.empty:
                  raise ValueError("table is empty"  )
           if query is None or query == "":
                  raise ValueError("query is empty"  )
           lowerCamelCase     :   Any    				=      self.tokenizer(UpperCamelCase__  ,						UpperCamelCase__  ,						return_tensors=self.framework  ,						truncation=UpperCamelCase__  ,						padding=UpperCamelCase__  )
           lowerCamelCase     :   str    				=      table
           return inputs







    def 					_lowercase  (					self  ,						UpperCamelCase__  ,						UpperCamelCase__=False  )  ->   Dict:
           lowerCamelCase     :   Union[str, Any]    				=      model_inputs.pop("table"  )

           if self.type == "tapas":
                  if sequential:
                         lowerCamelCase     :   int    				=      self.sequential_inference(**UpperCamelCase__  )
                  else:
                         lowerCamelCase     :   List[Any]    				=      self.batch_inference(**UpperCamelCase__  )
           else:
                  lowerCamelCase     :   int    				=      self.model.generate(**UpperCamelCase__  )
           lowerCamelCase     :   int    				=      {"model_inputs": model_inputs, "table": table, "outputs": outputs}
           return model_outputs







    def 					_lowercase  (					self  ,						UpperCamelCase__  )  ->   List[str]:
           lowerCamelCase     :   List[Any]    				=      model_outputs["model_inputs"]
           lowerCamelCase     :   Tuple    				=      model_outputs["table"]
           lowerCamelCase     :   Optional[Any]    				=      model_outputs["outputs"]
           if self.type == "tapas":
                  if self.aggregate:
                         lowerCamelCase				,   lowerCamelCase     :   Tuple    				=      outputs[:2]
                         lowerCamelCase     :   Union[str, Any]    				=      self.tokenizer.convert_logits_to_predictions(UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  )
                         lowerCamelCase				,   lowerCamelCase     :   Dict    				=      predictions
                         lowerCamelCase     :   int    				=      {i: self.model.config.aggregation_labels[pred] for i, pred in enumerate(UpperCamelCase__  )}

                         lowerCamelCase     :   Tuple    				=      self.model.config.no_aggregation_label_index
                         lowerCamelCase     :   int    				=      {
                             i: aggregators[i] + " > " for i, pred in enumerate(UpperCamelCase__  ) if pred != no_agg_label_index
                         }
                  else:
                         lowerCamelCase     :   Optional[int]    				=      outputs[0]
                         lowerCamelCase     :   Tuple    				=      self.tokenizer.convert_logits_to_predictions(UpperCamelCase__  ,						UpperCamelCase__  )
                         lowerCamelCase     :   Optional[Any]    				=      predictions[0]
                         lowerCamelCase     :   Dict    				=      {}
                         lowerCamelCase     :   str    				=      {}
                  lowerCamelCase     :   Optional[int]    				=      []
                  for index, coordinates in enumerate(UpperCamelCase__  ):
                         lowerCamelCase     :   Optional[Any]    				=      [table.iat[coordinate] for coordinate in coordinates]
                         lowerCamelCase     :   Union[str, Any]    				=      aggregators.get(UpperCamelCase__  ,						""  )
                         lowerCamelCase     :   str    				=      aggregators_prefix.get(UpperCamelCase__  ,						""  )
                         lowerCamelCase     :   Tuple    				=      {
                             "answer": aggregator_prefix + ", ".join(UpperCamelCase__  ),
                             "coordinates": coordinates,
                             "cells": [table.iat[coordinate] for coordinate in coordinates],
                         }
                         if aggregator:
                                lowerCamelCase     :   str    				=      aggregator

                         answers.append(UpperCamelCase__  )
                  if len(UpperCamelCase__  ) == 0:
                         raise PipelineException("Empty answer"  )
           else:
                  lowerCamelCase     :   Optional[Any]    				=      [{"answer": answer} for answer in self.tokenizer.batch_decode(UpperCamelCase__  ,						skip_special_tokens=UpperCamelCase__  )]

           return answers if len(UpperCamelCase__  ) > 1 else answers[0]






