import json
import os
from dataclasses import dataclass
from functools import partial
from typing import Callable

import flax.linen as nn
import jax
import jax.numpy as jnp
import joblib
import optax
import wandb
from flax import jax_utils, struct, traverse_util
from flax.serialization import from_bytes, to_bytes
from flax.training import train_state
from flax.training.common_utils import shard
from tqdm.auto import tqdm

from transformers import BigBirdConfig, FlaxBigBirdForQuestionAnswering
from transformers.models.big_bird.modeling_flax_big_bird import FlaxBigBirdForQuestionAnsweringModule


class      lowerCamelCase    (SCREAMING_SNAKE_CASE__    ):
  """simple docstring"""







  lowerCamelCase__   =  42
  lowerCamelCase__   =  jnp.floataa
  lowerCamelCase__   =  True







  def       __A   (							self       :				Union[str, Any]      )     ->		Optional[int]:
        super().setup()
        SCREAMING_SNAKE_CASE_      =     nn.Dense(5     ,   dtype=self.dtype      )







  def __call__(							self       :				int     ,   *__magic_name__       :				Optional[int]     ,   **__magic_name__       :				Optional[int]      )     ->		Union[str, Any]:
        SCREAMING_SNAKE_CASE_      =     super().__call__(*__magic_name__     ,   **__magic_name__      )
        SCREAMING_SNAKE_CASE_      =     self.cls(outputs[2]      )
        return outputs[:2] + (cls_out,)



class      lowerCamelCase    (SCREAMING_SNAKE_CASE__    ):
  """simple docstring"""







  lowerCamelCase__   =  FlaxBigBirdForNaturalQuestionsModule



def    a__       (						__UpperCamelCase   ,					__UpperCamelCase   ,					__UpperCamelCase   ,					__UpperCamelCase   ,					__UpperCamelCase   ,					__UpperCamelCase ):
      def cross_entropy(__UpperCamelCase   ,					__UpperCamelCase   ,					__UpperCamelCase=None ):
            SCREAMING_SNAKE_CASE_      =     logits.shape[-1]
            SCREAMING_SNAKE_CASE_      =     (labels[..., None] == jnp.arange(__UpperCamelCase )[None]).astype("f4" )
            SCREAMING_SNAKE_CASE_      =     jax.nn.log_softmax(__UpperCamelCase   ,					axis=-1 )
            SCREAMING_SNAKE_CASE_      =     -jnp.sum(labels * logits   ,					axis=-1 )
            if reduction is not None:
                  SCREAMING_SNAKE_CASE_      =     reduction(__UpperCamelCase )
            return loss

      SCREAMING_SNAKE_CASE_      =     partial(__UpperCamelCase   ,					reduction=jnp.mean )
      SCREAMING_SNAKE_CASE_      =     cross_entropy(__UpperCamelCase   ,					__UpperCamelCase )
      SCREAMING_SNAKE_CASE_      =     cross_entropy(__UpperCamelCase   ,					__UpperCamelCase )
      SCREAMING_SNAKE_CASE_      =     cross_entropy(__UpperCamelCase   ,					__UpperCamelCase )
      return (start_loss + end_loss + pooled_loss) / 3


@dataclass
class      lowerCamelCase    :
  """simple docstring"""







  lowerCamelCase__   =  "google/bigbird-roberta-base"
  lowerCamelCase__   =  3_0_0_0
  lowerCamelCase__   =  1_0_5_0_0

  lowerCamelCase__   =  1_2_8
  lowerCamelCase__   =  3

  lowerCamelCase__   =  1
  lowerCamelCase__   =  5

  # tx_args
  lowerCamelCase__   =  3E-5
  lowerCamelCase__   =  0.0
  lowerCamelCase__   =  2_0_0_0_0
  lowerCamelCase__   =  0.0095

  lowerCamelCase__   =  "bigbird-roberta-natural-questions"
  lowerCamelCase__   =  "training-expt"
  lowerCamelCase__   =  "data/nq-training.jsonl"
  lowerCamelCase__   =  "data/nq-validation.jsonl"







  def       __A   (							self       :				List[Any]      )     ->		Union[str, Any]:
        os.makedirs(self.base_dir     ,   exist_ok=__magic_name__      )
        SCREAMING_SNAKE_CASE_      =     os.path.join(self.base_dir     ,   self.save_dir      )
        SCREAMING_SNAKE_CASE_      =     self.batch_size_per_device * jax.device_count()



@dataclass
class      lowerCamelCase    :
  """simple docstring"""







  lowerCamelCase__   =  42
  lowerCamelCase__   =  4_0_9_6  # no dynamic padding on TPUs







  def __call__(							self       :				Optional[int]     ,   __magic_name__       :				Tuple      )     ->		Tuple:
        SCREAMING_SNAKE_CASE_      =     self.collate_fn(__magic_name__      )
        SCREAMING_SNAKE_CASE_      =     jax.tree_util.tree_map(__magic_name__     ,   __magic_name__      )
        return batch







  def       __A   (							self       :				str     ,   __magic_name__       :				Optional[int]      )     ->		Tuple:
        SCREAMING_SNAKE_CASE_   ,       SCREAMING_SNAKE_CASE_      =     self.fetch_inputs(features["input_ids"]      )
        SCREAMING_SNAKE_CASE_      =     {
            "input_ids": jnp.array(__magic_name__     ,   dtype=jnp.intaa      ),
            "attention_mask": jnp.array(__magic_name__     ,   dtype=jnp.intaa      ),
            "start_labels": jnp.array(features["start_token"]     ,   dtype=jnp.intaa      ),
            "end_labels": jnp.array(features["end_token"]     ,   dtype=jnp.intaa      ),
            "pooled_labels": jnp.array(features["category"]     ,   dtype=jnp.intaa      ),
        }
        return batch







  def       __A   (							self       :				List[str]     ,   __magic_name__       :				list      )     ->		Optional[Any]:
        SCREAMING_SNAKE_CASE_      =     [self._fetch_inputs(__magic_name__      ) for ids in input_ids]
        return zip(*__magic_name__      )







  def       __A   (							self       :				str     ,   __magic_name__       :				list      )     ->		Optional[int]:
        SCREAMING_SNAKE_CASE_      =     [1 for _ in range(len(__magic_name__      )      )]
        while len(__magic_name__      ) < self.max_length:
              input_ids.append(self.pad_id      )
              attention_mask.append(0      )
        return input_ids, attention_mask



def    a__       (						__UpperCamelCase   ,					__UpperCamelCase   ,					__UpperCamelCase=None ):
      if seed is not None:
            SCREAMING_SNAKE_CASE_      =     dataset.shuffle(seed=__UpperCamelCase )
      for i in range(len(__UpperCamelCase ) // batch_size ):
            SCREAMING_SNAKE_CASE_      =     dataset[i * batch_size : (i + 1) * batch_size]
            yield dict(__UpperCamelCase )


@partial(jax.pmap   ,					axis_name="batch" )
def    a__       (						__UpperCamelCase   ,					__UpperCamelCase   ,					**__UpperCamelCase ):
      def loss_fn(__UpperCamelCase ):
            SCREAMING_SNAKE_CASE_      =     model_inputs.pop("start_labels" )
            SCREAMING_SNAKE_CASE_      =     model_inputs.pop("end_labels" )
            SCREAMING_SNAKE_CASE_      =     model_inputs.pop("pooled_labels" )

            SCREAMING_SNAKE_CASE_      =     state.apply_fn(**__UpperCamelCase   ,					params=__UpperCamelCase   ,					dropout_rng=__UpperCamelCase   ,					train=__UpperCamelCase )
            SCREAMING_SNAKE_CASE_   ,       SCREAMING_SNAKE_CASE_   ,       SCREAMING_SNAKE_CASE_      =     outputs

            return state.loss_fn(
                __UpperCamelCase   ,					__UpperCamelCase   ,					__UpperCamelCase   ,					__UpperCamelCase   ,					__UpperCamelCase   ,					__UpperCamelCase   ,					)

      SCREAMING_SNAKE_CASE_   ,       SCREAMING_SNAKE_CASE_      =     jax.random.split(__UpperCamelCase )
      SCREAMING_SNAKE_CASE_      =     jax.value_and_grad(__UpperCamelCase )
      SCREAMING_SNAKE_CASE_   ,       SCREAMING_SNAKE_CASE_      =     grad_fn(state.params )
      SCREAMING_SNAKE_CASE_      =     jax.lax.pmean({"loss": loss}   ,					axis_name="batch" )
      SCREAMING_SNAKE_CASE_      =     jax.lax.pmean(__UpperCamelCase   ,					"batch" )

      SCREAMING_SNAKE_CASE_      =     state.apply_gradients(grads=__UpperCamelCase )
      return state, metrics, new_drp_rng


@partial(jax.pmap   ,					axis_name="batch" )
def    a__       (						__UpperCamelCase   ,					**__UpperCamelCase ):
      SCREAMING_SNAKE_CASE_      =     model_inputs.pop("start_labels" )
      SCREAMING_SNAKE_CASE_      =     model_inputs.pop("end_labels" )
      SCREAMING_SNAKE_CASE_      =     model_inputs.pop("pooled_labels" )

      SCREAMING_SNAKE_CASE_      =     state.apply_fn(**__UpperCamelCase   ,					params=state.params   ,					train=__UpperCamelCase )
      SCREAMING_SNAKE_CASE_   ,       SCREAMING_SNAKE_CASE_   ,       SCREAMING_SNAKE_CASE_      =     outputs

      SCREAMING_SNAKE_CASE_      =     state.loss_fn(__UpperCamelCase   ,					__UpperCamelCase   ,					__UpperCamelCase   ,					__UpperCamelCase   ,					__UpperCamelCase   ,					__UpperCamelCase )
      SCREAMING_SNAKE_CASE_      =     jax.lax.pmean({"loss": loss}   ,					axis_name="batch" )
      return metrics


class      lowerCamelCase    (train_state.TrainState    ):
  """simple docstring"""







  lowerCamelCase__   =  struct.field(pytree_node=SCREAMING_SNAKE_CASE__    )



@dataclass
class      lowerCamelCase    :
  """simple docstring"""







  lowerCamelCase__   =  42
  lowerCamelCase__   =  42
  lowerCamelCase__   =  42
  lowerCamelCase__   =  42
  lowerCamelCase__   =  42
  lowerCamelCase__   =  42
  lowerCamelCase__   =  None







  def       __A   (							self       :				str     ,   __magic_name__       :				Dict     ,   __magic_name__       :				str     ,   __magic_name__       :				Union[str, Any]     ,   __magic_name__       :				Union[str, Any]=None      )     ->		Dict:
        SCREAMING_SNAKE_CASE_      =     model.params
        SCREAMING_SNAKE_CASE_      =     TrainState.create(
            apply_fn=model.__call__     ,   params=__magic_name__     ,   tx=__magic_name__     ,   loss_fn=__magic_name__     ,   )
        if ckpt_dir is not None:
              SCREAMING_SNAKE_CASE_   ,       SCREAMING_SNAKE_CASE_   ,       SCREAMING_SNAKE_CASE_   ,       SCREAMING_SNAKE_CASE_   ,       SCREAMING_SNAKE_CASE_      =     restore_checkpoint(__magic_name__     ,   __magic_name__      )
              SCREAMING_SNAKE_CASE_      =     {
                  "lr": args.lr,
                  "init_lr": args.init_lr,
                  "warmup_steps": args.warmup_steps,
                  "num_train_steps": num_train_steps,
                  "weight_decay": args.weight_decay,
              }
              SCREAMING_SNAKE_CASE_   ,       SCREAMING_SNAKE_CASE_      =     build_tx(**__magic_name__      )
              SCREAMING_SNAKE_CASE_      =     train_state.TrainState(
                  step=__magic_name__     ,   apply_fn=model.__call__     ,   params=__magic_name__     ,   tx=__magic_name__     ,   opt_state=__magic_name__     ,   )
              SCREAMING_SNAKE_CASE_      =     args
              SCREAMING_SNAKE_CASE_      =     data_collator
              SCREAMING_SNAKE_CASE_      =     lr
              SCREAMING_SNAKE_CASE_      =     params
        SCREAMING_SNAKE_CASE_      =     jax_utils.replicate(__magic_name__      )
        return state







  def       __A   (							self       :				str     ,   __magic_name__       :				Optional[Any]     ,   __magic_name__       :				Optional[int]     ,   __magic_name__       :				str      )     ->		List[Any]:
        SCREAMING_SNAKE_CASE_      =     self.args
        SCREAMING_SNAKE_CASE_      =     len(__magic_name__      ) // args.batch_size

        SCREAMING_SNAKE_CASE_      =     jax.random.PRNGKey(0      )
        SCREAMING_SNAKE_CASE_      =     jax.random.split(__magic_name__     ,   jax.device_count()      )
        for epoch in range(args.max_epochs      ):
              SCREAMING_SNAKE_CASE_      =     jnp.array(0     ,   dtype=jnp.floataa      )
              SCREAMING_SNAKE_CASE_      =     get_batched_dataset(__magic_name__     ,   args.batch_size     ,   seed=__magic_name__      )
              SCREAMING_SNAKE_CASE_      =     0
              for batch in tqdm(__magic_name__     ,   total=__magic_name__     ,   desc=F'''Running EPOCH-{epoch}'''      ):
                    SCREAMING_SNAKE_CASE_      =     self.data_collator(__magic_name__      )
                    SCREAMING_SNAKE_CASE_   ,       SCREAMING_SNAKE_CASE_   ,       SCREAMING_SNAKE_CASE_      =     self.train_step_fn(__magic_name__     ,   __magic_name__     ,   **__magic_name__      )
                    running_loss += jax_utils.unreplicate(metrics["loss"]      )
                    i += 1
                    if i % args.logging_steps == 0:
                          SCREAMING_SNAKE_CASE_      =     jax_utils.unreplicate(state.step      )
                          SCREAMING_SNAKE_CASE_      =     running_loss.item() / i
                          SCREAMING_SNAKE_CASE_      =     self.scheduler_fn(state_step - 1      )

                          SCREAMING_SNAKE_CASE_      =     self.evaluate(__magic_name__     ,   __magic_name__      )
                          SCREAMING_SNAKE_CASE_      =     {
                              "step": state_step.item(),
                              "eval_loss": eval_loss.item(),
                              "tr_loss": tr_loss,
                              "lr": lr.item(),
                          }
                          tqdm.write(str(__magic_name__      )      )
                          self.logger.log(__magic_name__     ,   commit=__magic_name__      )

                    if i % args.save_steps == 0:
                          self.save_checkpoint(args.save_dir + F'''-e{epoch}-s{i}'''     ,   state=__magic_name__      )







  def       __A   (							self       :				Tuple     ,   __magic_name__       :				Union[str, Any]     ,   __magic_name__       :				Optional[int]      )     ->		Optional[Any]:
        SCREAMING_SNAKE_CASE_      =     get_batched_dataset(__magic_name__     ,   self.args.batch_size      )
        SCREAMING_SNAKE_CASE_      =     len(__magic_name__      ) // self.args.batch_size
        SCREAMING_SNAKE_CASE_      =     jnp.array(0     ,   dtype=jnp.floataa      )
        SCREAMING_SNAKE_CASE_      =     0
        for batch in tqdm(__magic_name__     ,   total=__magic_name__     ,   desc="Evaluating ... "      ):
              SCREAMING_SNAKE_CASE_      =     self.data_collator(__magic_name__      )
              SCREAMING_SNAKE_CASE_      =     self.val_step_fn(__magic_name__     ,   **__magic_name__      )
              running_loss += jax_utils.unreplicate(metrics["loss"]      )
              i += 1
        return running_loss / i







  def       __A   (							self       :				str     ,   __magic_name__       :				Optional[int]     ,   __magic_name__       :				int      )     ->		Optional[Any]:
        SCREAMING_SNAKE_CASE_      =     jax_utils.unreplicate(__magic_name__      )
        print(F'''SAVING CHECKPOINT IN {save_dir}'''     ,   end=" ... "      )
        self.model_save_fn(__magic_name__     ,   params=state.params      )
        with open(os.path.join(__magic_name__     ,   "opt_state.msgpack"      )     ,   "wb"      ) as f:
              f.write(to_bytes(state.opt_state      )      )
        joblib.dump(self.args     ,   os.path.join(__magic_name__     ,   "args.joblib"      )      )
        joblib.dump(self.data_collator     ,   os.path.join(__magic_name__     ,   "data_collator.joblib"      )      )
        with open(os.path.join(__magic_name__     ,   "training_state.json"      )     ,   "w"      ) as f:
              json.dump({"step": state.step.item()}     ,   __magic_name__      )
        print("DONE"      )



def    a__       (						__UpperCamelCase   ,					__UpperCamelCase ):
      print(F'''RESTORING CHECKPOINT FROM {save_dir}'''   ,					end=" ... " )
      with open(os.path.join(__UpperCamelCase   ,					"flax_model.msgpack" )   ,					"rb" ) as f:
            SCREAMING_SNAKE_CASE_      =     from_bytes(state.params   ,					f.read() )

      with open(os.path.join(__UpperCamelCase   ,					"opt_state.msgpack" )   ,					"rb" ) as f:
            SCREAMING_SNAKE_CASE_      =     from_bytes(state.opt_state   ,					f.read() )

      SCREAMING_SNAKE_CASE_      =     joblib.load(os.path.join(__UpperCamelCase   ,					"args.joblib" ) )
      SCREAMING_SNAKE_CASE_      =     joblib.load(os.path.join(__UpperCamelCase   ,					"data_collator.joblib" ) )

      with open(os.path.join(__UpperCamelCase   ,					"training_state.json" )   ,					"r" ) as f:
            SCREAMING_SNAKE_CASE_      =     json.load(__UpperCamelCase )
      SCREAMING_SNAKE_CASE_      =     training_state["step"]

      print("DONE" )
      return params, opt_state, step, args, data_collator


def    a__       (						__UpperCamelCase   ,					__UpperCamelCase   ,					__UpperCamelCase   ,					__UpperCamelCase ):
      SCREAMING_SNAKE_CASE_      =     num_train_steps - warmup_steps
      SCREAMING_SNAKE_CASE_      =     optax.linear_schedule(init_value=__UpperCamelCase   ,					end_value=__UpperCamelCase   ,					transition_steps=__UpperCamelCase )
      SCREAMING_SNAKE_CASE_      =     optax.linear_schedule(init_value=__UpperCamelCase   ,					end_value=1E-7   ,					transition_steps=__UpperCamelCase )
      SCREAMING_SNAKE_CASE_      =     optax.join_schedules(schedules=[warmup_fn, decay_fn]   ,					boundaries=[warmup_steps] )
      return lr





def    a__       (						__UpperCamelCase   ,					__UpperCamelCase   ,					__UpperCamelCase   ,					__UpperCamelCase   ,					__UpperCamelCase ):
      def weight_decay_mask(__UpperCamelCase ):
            SCREAMING_SNAKE_CASE_      =     traverse_util.flatten_dict(__UpperCamelCase )
            SCREAMING_SNAKE_CASE_      =     {k: (v[-1] != "bias" and v[-2:] != ("LayerNorm", "scale")) for k, v in params.items()}
            return traverse_util.unflatten_dict(__UpperCamelCase )

      SCREAMING_SNAKE_CASE_      =     scheduler_fn(__UpperCamelCase   ,					__UpperCamelCase   ,					__UpperCamelCase   ,					__UpperCamelCase )

      SCREAMING_SNAKE_CASE_      =     optax.adamw(learning_rate=__UpperCamelCase   ,					weight_decay=__UpperCamelCase   ,					mask=__UpperCamelCase )
      return tx, lr



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


A   :     Tuple       =				get_tests_dir("fixtures")


class      lowerCamelCase    (unittest.TestCase    ):
  """simple docstring"""







  def       __A   (							self       :				Union[str, Any]      )     ->		Union[str, Any]:
        # A mock response for an HTTP head request to emulate server down
        SCREAMING_SNAKE_CASE_      =     mock.Mock()
        SCREAMING_SNAKE_CASE_      =     500
        SCREAMING_SNAKE_CASE_      =     {}
        SCREAMING_SNAKE_CASE_      =     HTTPError
        SCREAMING_SNAKE_CASE_      =     {}

        # Download this model to make sure it's in the cache.
        SCREAMING_SNAKE_CASE_      =     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=__magic_name__      ) as mock_head:
              SCREAMING_SNAKE_CASE_      =     WavaVecaFeatureExtractor.from_pretrained("hf-internal-testing/tiny-random-wav2vec2"      )
              # This check we did call the fake head request
              mock_head.assert_called()







  def       __A   (							self       :				Optional[int]      )     ->		Tuple:
        # This test is for deprecated behavior and can be removed in v5
        SCREAMING_SNAKE_CASE_      =     WavaVecaFeatureExtractor.from_pretrained(
            "https://huggingface.co/hf-internal-testing/tiny-random-wav2vec2/resolve/main/preprocessor_config.json"      )



@is_staging_test
class      lowerCamelCase    (unittest.TestCase    ):
  """simple docstring"""







  @classmethod
  def       __A   (							cls       :				Any      )     ->		Union[str, Any]:
        SCREAMING_SNAKE_CASE_      =     TOKEN
        HfFolder.save_token(__magic_name__      )







  @classmethod
  def       __A   (							cls       :				int      )     ->		Optional[Any]:
        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       __A   (							self       :				str      )     ->		List[Any]:
        SCREAMING_SNAKE_CASE_      =     WavaVecaFeatureExtractor.from_pretrained(__magic_name__      )
        feature_extractor.push_to_hub("test-feature-extractor"     ,   use_auth_token=self._token      )

        SCREAMING_SNAKE_CASE_      =     WavaVecaFeatureExtractor.from_pretrained(F'''{USER}/test-feature-extractor'''      )
        for k, v in feature_extractor.__dict__.items():
              self.assertEqual(__magic_name__     ,   getattr(__magic_name__     ,   __magic_name__      )      )

        # 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(
                  __magic_name__     ,   repo_id="test-feature-extractor"     ,   push_to_hub=__magic_name__     ,   use_auth_token=self._token      )

        SCREAMING_SNAKE_CASE_      =     WavaVecaFeatureExtractor.from_pretrained(F'''{USER}/test-feature-extractor'''      )
        for k, v in feature_extractor.__dict__.items():
              self.assertEqual(__magic_name__     ,   getattr(__magic_name__     ,   __magic_name__      )      )







  def       __A   (							self       :				Dict      )     ->		str:
        SCREAMING_SNAKE_CASE_      =     WavaVecaFeatureExtractor.from_pretrained(__magic_name__      )
        feature_extractor.push_to_hub("valid_org/test-feature-extractor"     ,   use_auth_token=self._token      )

        SCREAMING_SNAKE_CASE_      =     WavaVecaFeatureExtractor.from_pretrained("valid_org/test-feature-extractor"      )
        for k, v in feature_extractor.__dict__.items():
              self.assertEqual(__magic_name__     ,   getattr(__magic_name__     ,   __magic_name__      )      )

        # 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(
                  __magic_name__     ,   repo_id="valid_org/test-feature-extractor-org"     ,   push_to_hub=__magic_name__     ,   use_auth_token=self._token      )

        SCREAMING_SNAKE_CASE_      =     WavaVecaFeatureExtractor.from_pretrained("valid_org/test-feature-extractor-org"      )
        for k, v in feature_extractor.__dict__.items():
              self.assertEqual(__magic_name__     ,   getattr(__magic_name__     ,   __magic_name__      )      )







  def       __A   (							self       :				List[Any]      )     ->		Dict:
        CustomFeatureExtractor.register_for_auto_class()
        SCREAMING_SNAKE_CASE_      =     CustomFeatureExtractor.from_pretrained(__magic_name__      )

        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_      =     AutoFeatureExtractor.from_pretrained(
            F'''{USER}/test-dynamic-feature-extractor'''     ,   trust_remote_code=__magic_name__      )
        # 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"      )


