

import numpy as np
import torch
from torch.utils.data import Dataset, IterableDataset

from ..utils.generic import ModelOutput







class        a__       (   UpperCAmelCase       ):







  """simple docstring"""




  def __init__(			self    : Dict				,				UpperCAmelCase__    : Optional[Any]				,				UpperCAmelCase__    : int				,				UpperCAmelCase__    : str		)					->List[str]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Union[str, Any]       =   dataset
        SCREAMING_SNAKE_CASE :      List[str]       =   process
        SCREAMING_SNAKE_CASE :      List[str]       =   params


  def __len__(			self    : Dict		)					->List[Any]:
        """simple docstring"""







        return len(self.dataset		)




  def __getitem__(			self    : Any				,				UpperCAmelCase__    : Optional[int]		)					->Optional[int]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      List[Any]       =   self.dataset[i]
        SCREAMING_SNAKE_CASE :      str       =   self.process(UpperCAmelCase__				,				**self.params		)
        return processed







class        a__       (   UpperCAmelCase       ):







  """simple docstring"""




  def __init__(			self    : Union[str, Any]				,				UpperCAmelCase__    : Tuple				,				UpperCAmelCase__    : List[str]				,				UpperCAmelCase__    : Tuple				,				UpperCAmelCase__    : Tuple=None		)					->Dict:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      List[Any]       =   loader
        SCREAMING_SNAKE_CASE :      List[Any]       =   infer
        SCREAMING_SNAKE_CASE :      Union[str, Any]       =   params
        if loader_batch_size == 1:
              # Let's spare some time by deactivating altogether
              SCREAMING_SNAKE_CASE :      List[Any]       =   None
        SCREAMING_SNAKE_CASE :      Any       =   loader_batch_size

        # Internal bookkeeping
        SCREAMING_SNAKE_CASE :      Optional[int]       =   None
        SCREAMING_SNAKE_CASE :      Tuple       =   None


  def __len__(			self    : str		)					->Optional[int]:
        """simple docstring"""







        return len(self.loader		)


  def __iter__(			self    : Dict		)					->Union[str, Any]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      str       =   iter(self.loader		)
        return self


  def  _lowercase      (			self    : Optional[int]		)					->Union[str, Any]:
        """simple docstring"""







        if isinstance(self._loader_batch_data				,				torch.Tensor		):
              # Batch data is simple tensor, just fetch the slice
              SCREAMING_SNAKE_CASE :      Optional[int]       =   self._loader_batch_data[self._loader_batch_index]
        else:
              # Batch data is assumed to be BaseModelOutput (or dict)
              SCREAMING_SNAKE_CASE :      List[Any]       =   {}
              for k, element in self._loader_batch_data.items():
                    if isinstance(UpperCAmelCase__				,				UpperCAmelCase__		):
                          # Convert ModelOutput to tuple first
                          SCREAMING_SNAKE_CASE :      Tuple       =   element.to_tuple()
                          if isinstance(element[0]				,				torch.Tensor		):
                                SCREAMING_SNAKE_CASE :      Optional[int]       =   tuple(el[self._loader_batch_index].unsqueeze(0		) for el in element		)
                          elif isinstance(element[0]				,				np.ndarray		):
                                SCREAMING_SNAKE_CASE :      Tuple       =   tuple(np.expand_dims(el[self._loader_batch_index]				,				0		) for el in element		)
                          continue
                    if k in {"hidden_states", "past_key_values", "attentions"} and isinstance(UpperCAmelCase__				,				UpperCAmelCase__		):
                          # Those are stored as lists of tensors so need specific unbatching.
                          if isinstance(element[0]				,				torch.Tensor		):
                                SCREAMING_SNAKE_CASE :      Optional[Any]       =   tuple(el[self._loader_batch_index].unsqueeze(0		) for el in element		)
                          elif isinstance(element[0]				,				np.ndarray		):
                                SCREAMING_SNAKE_CASE :      Tuple       =   tuple(np.expand_dims(el[self._loader_batch_index]				,				0		) for el in element		)
                          continue
                    if element is None:
                          # This can happen for optional data that get passed around
                          SCREAMING_SNAKE_CASE :      Optional[Any]       =   None
                    elif isinstance(element[self._loader_batch_index]				,				torch.Tensor		):
                          # Take correct batch data, but make it looked like batch_size=1
                          # For compatibility with other methods within transformers

                          SCREAMING_SNAKE_CASE :      List[str]       =   element[self._loader_batch_index].unsqueeze(0		)
                    elif isinstance(element[self._loader_batch_index]				,				np.ndarray		):
                          # Take correct batch data, but make it looked like batch_size=1
                          # For compatibility with other methods within transformers
                          SCREAMING_SNAKE_CASE :      int       =   np.expand_dims(element[self._loader_batch_index]				,				0		)
                    else:
                          # This is typically a list, so no need to `unsqueeze`.
                          SCREAMING_SNAKE_CASE :      Optional[int]       =   element[self._loader_batch_index]
            # Recreate the element by reusing the original class to make it look
            # batch_size=1
              SCREAMING_SNAKE_CASE :      str       =   self._loader_batch_data.__class__(UpperCAmelCase__		)
        self._loader_batch_index += 1
        return result




  def  _lowercase      (			self    : Any		)					->int:
        """simple docstring"""







        if self._loader_batch_index is not None and self._loader_batch_index < self.loader_batch_size:
              # We are currently unrolling a batch so we just need to return
              # the current item within a batch
              return self.loader_batch_item()

        # We're out of items within a batch
        SCREAMING_SNAKE_CASE :      List[str]       =   next(self.iterator		)
        SCREAMING_SNAKE_CASE :      Any       =   self.infer(UpperCAmelCase__				,				**self.params		)
        # We now have a batch of "inferred things".
        if self.loader_batch_size is not None:
              # Try to infer the size of the batch
              if isinstance(UpperCAmelCase__				,				torch.Tensor		):
                    SCREAMING_SNAKE_CASE :      List[str]       =   processed
              else:
                    SCREAMING_SNAKE_CASE :      Union[str, Any]       =   list(processed.keys()		)[0]
                    SCREAMING_SNAKE_CASE :      Tuple       =   processed[key]
              if isinstance(UpperCAmelCase__				,				UpperCAmelCase__		):
                    SCREAMING_SNAKE_CASE :      str       =   len(UpperCAmelCase__		)
              else:
                    SCREAMING_SNAKE_CASE :      Optional[Any]       =   first_tensor.shape[0]
              if 0 < observed_batch_size < self.loader_batch_size:
                    # could be last batch so we can't unroll as many
                    # elements.
                    SCREAMING_SNAKE_CASE :      Optional[Any]       =   observed_batch_size
              # Setting internal index to unwrap the batch
              SCREAMING_SNAKE_CASE :      Optional[Any]       =   processed
              SCREAMING_SNAKE_CASE :      Union[str, Any]       =   0
              return self.loader_batch_item()
        else:
              # We're not unrolling batches
              return processed







class        a__       (   UpperCAmelCase       ):







  """simple docstring"""




  def __init__(			self    : Tuple				,				UpperCAmelCase__    : List[Any]				,				UpperCAmelCase__    : Optional[int]				,				UpperCAmelCase__    : List[str]				,				UpperCAmelCase__    : Tuple=None		)					->Union[str, Any]:
        """simple docstring"""







        super().__init__(UpperCAmelCase__				,				UpperCAmelCase__				,				UpperCAmelCase__		)


  def __iter__(			self    : Union[str, Any]		)					->List[Any]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      str       =   iter(self.loader		)
        SCREAMING_SNAKE_CASE :      Any       =   None
        return self




  def  _lowercase      (			self    : Optional[int]		)					->Dict:
        """simple docstring"""







        if self.subiterator is None:
              SCREAMING_SNAKE_CASE :      int       =   self.infer(next(self.iterator		)				,				**self.params		)
        try:
              # Try to return next item
              SCREAMING_SNAKE_CASE :      List[str]       =   next(self.subiterator		)
        except StopIteration:
              # When a preprocess iterator ends, we can start lookig at the next item
              # ChunkIterator will keep feeding until ALL elements of iterator
              # all have created their subiterator and have been iterating against.
              #
              # Another way to look at it, is we're basically flattening lists of lists
              # into a single list, but with generators
              SCREAMING_SNAKE_CASE :      List[Any]       =   self.infer(next(self.iterator		)				,				**self.params		)
              SCREAMING_SNAKE_CASE :      int       =   next(self.subiterator		)
        return processed







class        a__       (   UpperCAmelCase       ):







  """simple docstring"""




  def __iter__(			self    : Any		)					->Optional[int]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Any       =   iter(self.loader		)
        return self




  def  _lowercase      (			self    : List[Any]		)					->List[Any]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      int       =   False
        SCREAMING_SNAKE_CASE :      List[Any]       =   []
        if self._loader_batch_index is not None and self._loader_batch_index < self.loader_batch_size:
              while self._loader_batch_index < self.loader_batch_size:
                    SCREAMING_SNAKE_CASE :      int       =   self.loader_batch_item()
                    SCREAMING_SNAKE_CASE :      Optional[Any]       =   item.pop("""is_last"""		)
                    accumulator.append(UpperCAmelCase__		)
                    if is_last:
                          return accumulator

        while not is_last:
              SCREAMING_SNAKE_CASE :      Optional[Any]       =   self.infer(next(self.iterator		)				,				**self.params		)
              if self.loader_batch_size is not None:
                    if isinstance(UpperCAmelCase__				,				torch.Tensor		):
                          SCREAMING_SNAKE_CASE :      Optional[Any]       =   processed
                    else:
                          SCREAMING_SNAKE_CASE :      int       =   list(processed.keys()		)[0]
                          SCREAMING_SNAKE_CASE :      Any       =   processed[key]
                    if isinstance(UpperCAmelCase__				,				UpperCAmelCase__		):
                          SCREAMING_SNAKE_CASE :      int       =   len(UpperCAmelCase__		)
                    else:
                          SCREAMING_SNAKE_CASE :      Tuple       =   first_tensor.shape[0]
                    if 0 < observed_batch_size < self.loader_batch_size:
                          # could be last batch so we can't unroll as many
                          # elements.
                          SCREAMING_SNAKE_CASE :      List[Any]       =   observed_batch_size
                    SCREAMING_SNAKE_CASE :      Optional[int]       =   processed
                    SCREAMING_SNAKE_CASE :      Tuple       =   0
                    while self._loader_batch_index < self.loader_batch_size:
                          SCREAMING_SNAKE_CASE :      Optional[int]       =   self.loader_batch_item()
                          SCREAMING_SNAKE_CASE :      int       =   item.pop("""is_last"""		)
                          accumulator.append(UpperCAmelCase__		)
                          if is_last:
                                return accumulator
              else:
                    SCREAMING_SNAKE_CASE :      str       =   processed
                    SCREAMING_SNAKE_CASE :      Union[str, Any]       =   item.pop("""is_last"""		)
                    accumulator.append(UpperCAmelCase__		)
        return accumulator







class        a__       (   UpperCAmelCase       ):







  """simple docstring"""




  def __init__(			self    : Optional[int]				,				UpperCAmelCase__    : Dataset				,				UpperCAmelCase__    : str		)					->Optional[int]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Any       =   dataset
        SCREAMING_SNAKE_CASE :      Tuple       =   key


  def __len__(			self    : List[str]		)					->Tuple:
        """simple docstring"""







        return len(self.dataset		)




  def __getitem__(			self    : int				,				UpperCAmelCase__    : Optional[int]		)					->Dict:
        """simple docstring"""







        return self.dataset[i][self.key]







class        a__       (   UpperCAmelCase       ):







  """simple docstring"""




  def __init__(			self    : Dict				,				UpperCAmelCase__    : Dataset				,				UpperCAmelCase__    : str				,				UpperCAmelCase__    : str		)					->int:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Tuple       =   dataset
        SCREAMING_SNAKE_CASE :      Optional[int]       =   keya
        SCREAMING_SNAKE_CASE :      List[Any]       =   keya


  def __len__(			self    : str		)					->Optional[Any]:
        """simple docstring"""







        return len(self.dataset		)




  def __getitem__(			self    : Union[str, Any]				,				UpperCAmelCase__    : str		)					->int:
        """simple docstring"""







        return {"text": self.dataset[i][self.keya], "text_pair": self.dataset[i][self.keya]}







import sys
import tempfile
import unittest
import unittest.mock as mock
from pathlib import Path

from huggingface_hub import HfFolder, delete_repo
from requests.exceptions import HTTPError

from transformers import AutoFeatureExtractor, WavaVecaFeatureExtractor
from transformers.testing_utils import TOKEN, USER, get_tests_dir, is_staging_test


sys.path.append(str(Path(__file__).parent.parent / """utils"""))

from test_module.custom_feature_extraction import CustomFeatureExtractor  # noqa E402


UpperCAmelCase__ :     int				    =    get_tests_dir("""fixtures""")







class        a__       (   unittest.TestCase       ):







  """simple docstring"""




  def  _lowercase      (			self    : Dict		)					->Any:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Dict       =   mock.Mock()
        SCREAMING_SNAKE_CASE :      Optional[int]       =   5_0_0
        SCREAMING_SNAKE_CASE :      List[str]       =   {}
        SCREAMING_SNAKE_CASE :      List[Any]       =   HTTPError
        SCREAMING_SNAKE_CASE :      int       =   {}

        # Download this model to make sure it's in the cache.
        SCREAMING_SNAKE_CASE :      List[Any]       =   WavaVecaFeatureExtractor.from_pretrained("""hf-internal-testing/tiny-random-wav2vec2"""		)
        # Under the mock environment we get a 500 error when trying to reach the model.
        with mock.patch("""requests.Session.request"""				,				return_value=UpperCAmelCase__		) as mock_head:
              SCREAMING_SNAKE_CASE :      Optional[int]       =   WavaVecaFeatureExtractor.from_pretrained("""hf-internal-testing/tiny-random-wav2vec2"""		)
              # This check we did call the fake head request
              mock_head.assert_called()




  def  _lowercase      (			self    : Optional[Any]		)					->int:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      str       =   WavaVecaFeatureExtractor.from_pretrained(
            """https://huggingface.co/hf-internal-testing/tiny-random-wav2vec2/resolve/main/preprocessor_config.json"""		)







@is_staging_test
class        a__       (   unittest.TestCase       ):







  """simple docstring"""




  @classmethod
  def  _lowercase      (			cls    : List[str]		)					->Dict:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      int       =   TOKEN
        HfFolder.save_token(UpperCAmelCase__		)


  @classmethod
  def  _lowercase      (			cls    : List[str]		)					->Optional[int]:
        """simple docstring"""







        try:
              delete_repo(token=cls._token				,				repo_id="""test-feature-extractor"""		)
        except HTTPError:
              pass

        try:
              delete_repo(token=cls._token				,				repo_id="""valid_org/test-feature-extractor-org"""		)
        except HTTPError:
              pass

        try:
              delete_repo(token=cls._token				,				repo_id="""test-dynamic-feature-extractor"""		)
        except HTTPError:
              pass


  def  _lowercase      (			self    : int		)					->Optional[int]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      int       =   WavaVecaFeatureExtractor.from_pretrained(UpperCAmelCase__		)
        feature_extractor.push_to_hub("""test-feature-extractor"""				,				use_auth_token=self._token		)

        SCREAMING_SNAKE_CASE :      Any       =   WavaVecaFeatureExtractor.from_pretrained(f"{USER}/test-feature-extractor"		)
        for k, v in feature_extractor.__dict__.items():
              self.assertEqual(UpperCAmelCase__				,				getattr(UpperCAmelCase__				,				UpperCAmelCase__		)		)

        # Reset repo
        delete_repo(token=self._token				,				repo_id="""test-feature-extractor"""		)

        # Push to hub via save_pretrained
        with tempfile.TemporaryDirectory() as tmp_dir:
              feature_extractor.save_pretrained(
                  UpperCAmelCase__				,				repo_id="""test-feature-extractor"""				,				push_to_hub=UpperCAmelCase__				,				use_auth_token=self._token		)

        SCREAMING_SNAKE_CASE :      List[str]       =   WavaVecaFeatureExtractor.from_pretrained(f"{USER}/test-feature-extractor"		)
        for k, v in feature_extractor.__dict__.items():
              self.assertEqual(UpperCAmelCase__				,				getattr(UpperCAmelCase__				,				UpperCAmelCase__		)		)


  def  _lowercase      (			self    : List[Any]		)					->List[Any]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      List[str]       =   WavaVecaFeatureExtractor.from_pretrained(UpperCAmelCase__		)
        feature_extractor.push_to_hub("""valid_org/test-feature-extractor"""				,				use_auth_token=self._token		)

        SCREAMING_SNAKE_CASE :      Optional[Any]       =   WavaVecaFeatureExtractor.from_pretrained("""valid_org/test-feature-extractor"""		)
        for k, v in feature_extractor.__dict__.items():
              self.assertEqual(UpperCAmelCase__				,				getattr(UpperCAmelCase__				,				UpperCAmelCase__		)		)

        # Reset repo
        delete_repo(token=self._token				,				repo_id="""valid_org/test-feature-extractor"""		)

        # Push to hub via save_pretrained
        with tempfile.TemporaryDirectory() as tmp_dir:
              feature_extractor.save_pretrained(
                  UpperCAmelCase__				,				repo_id="""valid_org/test-feature-extractor-org"""				,				push_to_hub=UpperCAmelCase__				,				use_auth_token=self._token		)

        SCREAMING_SNAKE_CASE :      List[Any]       =   WavaVecaFeatureExtractor.from_pretrained("""valid_org/test-feature-extractor-org"""		)
        for k, v in feature_extractor.__dict__.items():
              self.assertEqual(UpperCAmelCase__				,				getattr(UpperCAmelCase__				,				UpperCAmelCase__		)		)




  def  _lowercase      (			self    : Optional[Any]		)					->Optional[int]:
        """simple docstring"""







        CustomFeatureExtractor.register_for_auto_class()
        SCREAMING_SNAKE_CASE :      Any       =   CustomFeatureExtractor.from_pretrained(UpperCAmelCase__		)

        feature_extractor.push_to_hub("""test-dynamic-feature-extractor"""				,				use_auth_token=self._token		)

        # This has added the proper auto_map field to the config
        self.assertDictEqual(
            feature_extractor.auto_map				,				{"""AutoFeatureExtractor""": """custom_feature_extraction.CustomFeatureExtractor"""}				,				)

        SCREAMING_SNAKE_CASE :      Any       =   AutoFeatureExtractor.from_pretrained(
            f"{USER}/test-dynamic-feature-extractor"				,				trust_remote_code=UpperCAmelCase__		)
        # Can't make an isinstance check because the new_feature_extractor is from the CustomFeatureExtractor class of a dynamic module
        self.assertEqual(new_feature_extractor.__class__.__name__				,				"""CustomFeatureExtractor"""		)




