






"""simple docstring"""


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 						snake_case  (  UpperCAmelCase	):
    __magic_name__   					=					42
    __magic_name__   					=					jnp.floataa
    __magic_name__   					=					True





    def     lowerCamelCase__				(							self     : Optional[Any]      ):


      '''simple docstring'''




      super().setup()
      a					:      str    =       nn.Dense(5    ,		dtype=self.dtype      )





    def __call__(							self     : Any    ,		*A     : Any    ,		**A     : Optional[int]      ):


      '''simple docstring'''




      a					:      List[str]    =       super().__call__(*A    ,		**A      )
      a					:      Optional[Any]    =       self.cls(outputs[2]      )
      return outputs[:2] + (cls_out,)






class 						snake_case  (  UpperCAmelCase	):
    __magic_name__   					=					FlaxBigBirdForNaturalQuestionsModule




def       snake_case  (A_				:Any     ,      A_				:Any     ,      A_				:Dict     ,      A_				:int     ,      A_				:Optional[Any]     ,      A_				:List[str]				):







  '''simple docstring'''

  def cross_entropy(A_				:List[str]     ,      A_				:Optional[Any]     ,      A_				:List[str]=None				):
    a					:      Optional[Any]    =       logits.shape[-1]
    a					:      List[str]    =       (labels[..., None] == jnp.arange(A_				)[None]).astype('f4'				)
    a					:      List[Any]    =       jax.nn.log_softmax(A_     ,      axis=-1				)
    a					:      Union[str, Any]    =       -jnp.sum(labels * logits     ,      axis=-1				)
    if reduction is not None:
      a					:      Tuple    =       reduction(A_				)
    return loss

  a					:      List[Any]    =       partial(A_     ,      reduction=jnp.mean				)
  a					:      str    =       cross_entropy(A_     ,      A_				)
  a					:      List[Any]    =       cross_entropy(A_     ,      A_				)
  a					:      List[str]    =       cross_entropy(A_     ,      A_				)
  return (start_loss + end_loss + pooled_loss) / 3






@dataclass
class 						snake_case  :
    __magic_name__   					=					"google/bigbird-roberta-base"
    __magic_name__   					=					3000
    __magic_name__   					=					1_0500

    __magic_name__   					=					128
    __magic_name__   					=					3

    __magic_name__   					=					1
    __magic_name__   					=					5

    # tx_args
    __magic_name__   					=					3E-5
    __magic_name__   					=					0.0
    __magic_name__   					=					2_0000
    __magic_name__   					=					0.0095

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





    def     lowerCamelCase__				(							self     : List[Any]      ):


      '''simple docstring'''




      os.makedirs(self.base_dir    ,		exist_ok=A      )
      a					:      Dict    =       os.path.join(self.base_dir    ,		self.save_dir      )
      a					:      Optional[int]    =       self.batch_size_per_device * jax.device_count()






@dataclass
class 						snake_case  :
    __magic_name__   					=					42
    __magic_name__   					=					4096  # no dynamic padding on TPUs





    def __call__(							self     : List[Any]    ,		A     : Any      ):


      '''simple docstring'''




      a					:      List[Any]    =       self.collate_fn(A      )
      a					:      Dict    =       jax.tree_util.tree_map(A    ,		A      )
      return batch





    def     lowerCamelCase__				(							self     : Optional[int]    ,		A     : Optional[int]      ):


      '''simple docstring'''




      a,					a					:      Dict    =       self.fetch_inputs(features['input_ids']      )
      a					:      str    =       {
          'input_ids': jnp.array(A    ,		dtype=jnp.intaa      ),
          'attention_mask': jnp.array(A    ,		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     lowerCamelCase__				(							self     : int    ,		A     : list      ):


      '''simple docstring'''




      a					:      Tuple    =       [self._fetch_inputs(A      ) for ids in input_ids]
      return zip(*A      )





    def     lowerCamelCase__				(							self     : int    ,		A     : list      ):


      '''simple docstring'''




      a					:      Dict    =       [1 for _ in range(len(A      )      )]
      while len(A      ) < self.max_length:
        input_ids.append(self.pad_id      )
        attention_mask.append(0      )
      return input_ids, attention_mask




def       snake_case  (A_				:List[Any]     ,      A_				:Dict     ,      A_				:Optional[Any]=None				):







  '''simple docstring'''

  if seed is not None:
    a					:      Any    =       dataset.shuffle(seed=A_				)
  for i in range(len(A_				) // batch_size				):
    a					:      Optional[Any]    =       dataset[i * batch_size : (i + 1) * batch_size]
    yield dict(A_				)




@partial(jax.pmap     ,      axis_name='batch'				)
def       snake_case  (A_				:int     ,      A_				:int     ,      **A_				:Any				):







  '''simple docstring'''

  def loss_fn(A_				:Optional[Any]				):
    a					:      Tuple    =       model_inputs.pop('start_labels'				)
    a					:      Tuple    =       model_inputs.pop('end_labels'				)
    a					:      Tuple    =       model_inputs.pop('pooled_labels'				)

    a					:      List[Any]    =       state.apply_fn(**A_     ,      params=A_     ,      dropout_rng=A_     ,      train=A_				)
    a,					a,					a					:      Union[str, Any]    =       outputs

    return state.loss_fn(
        A_     ,      A_     ,      A_     ,      A_     ,      A_     ,      A_     ,      )

  a,					a					:      Any    =       jax.random.split(A_				)
  a					:      int    =       jax.value_and_grad(A_				)
  a,					a					:      List[Any]    =       grad_fn(state.params				)
  a					:      Optional[int]    =       jax.lax.pmean({'loss': loss}     ,      axis_name='batch'				)
  a					:      Optional[int]    =       jax.lax.pmean(A_     ,      'batch'				)

  a					:      Union[str, Any]    =       state.apply_gradients(grads=A_				)
  return state, metrics, new_drp_rng




@partial(jax.pmap     ,      axis_name='batch'				)
def       snake_case  (A_				:List[str]     ,      **A_				:str				):







  '''simple docstring'''

  a					:      Tuple    =       model_inputs.pop('start_labels'				)
  a					:      List[Any]    =       model_inputs.pop('end_labels'				)
  a					:      Dict    =       model_inputs.pop('pooled_labels'				)

  a					:      Dict    =       state.apply_fn(**A_     ,      params=state.params     ,      train=A_				)
  a,					a,					a					:      Any    =       outputs

  a					:      Dict    =       state.loss_fn(A_     ,      A_     ,      A_     ,      A_     ,      A_     ,      A_				)
  a					:      Optional[int]    =       jax.lax.pmean({'loss': loss}     ,      axis_name='batch'				)
  return metrics






class 						snake_case  (  train_state.TrainState	):
    __magic_name__   					=					struct.field(pytree_node=UpperCAmelCase	)






@dataclass
class 						snake_case  :
    __magic_name__   					=					42
    __magic_name__   					=					42
    __magic_name__   					=					42
    __magic_name__   					=					42
    __magic_name__   					=					42
    __magic_name__   					=					42
    __magic_name__   					=					None





    def     lowerCamelCase__				(							self     : Any    ,		A     : Optional[Any]    ,		A     : Optional[Any]    ,		A     : Tuple    ,		A     : Optional[Any]=None      ):


      '''simple docstring'''




      a					:      Union[str, Any]    =       model.params
      a					:      Tuple    =       TrainState.create(
          apply_fn=model.__call__    ,		params=A    ,		tx=A    ,		loss_fn=A    ,		)
      if ckpt_dir is not None:
        a,					a,					a,					a,					a					:      Tuple    =       restore_checkpoint(A    ,		A      )
        a					:      Union[str, Any]    =       {
            'lr': args.lr,
            'init_lr': args.init_lr,
            'warmup_steps': args.warmup_steps,
            'num_train_steps': num_train_steps,
            'weight_decay': args.weight_decay,
        }
        a,					a					:      Any    =       build_tx(**A      )
        a					:      Any    =       train_state.TrainState(
            step=A    ,		apply_fn=model.__call__    ,		params=A    ,		tx=A    ,		opt_state=A    ,		)
        a					:      Tuple    =       args
        a					:      List[str]    =       data_collator
        a					:      Optional[Any]    =       lr
        a					:      str    =       params
      a					:      Tuple    =       jax_utils.replicate(A      )
      return state





    def     lowerCamelCase__				(							self     : List[str]    ,		A     : List[str]    ,		A     : Dict    ,		A     : Union[str, Any]      ):


      '''simple docstring'''




      a					:      int    =       self.args
      a					:      Optional[Any]    =       len(A      ) // args.batch_size

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

            a					:      Optional[Any]    =       self.evaluate(A    ,		A      )
            a					:      List[str]    =       {
                'step': state_step.item(),
                'eval_loss': eval_loss.item(),
                'tr_loss': tr_loss,
                'lr': lr.item(),
            }
            tqdm.write(str(A      )      )
            self.logger.log(A    ,		commit=A      )

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





    def     lowerCamelCase__				(							self     : List[Any]    ,		A     : Optional[int]    ,		A     : Dict      ):


      '''simple docstring'''




      a					:      Any    =       get_batched_dataset(A    ,		self.args.batch_size      )
      a					:      Dict    =       len(A      ) // self.args.batch_size
      a					:      Union[str, Any]    =       jnp.array(0    ,		dtype=jnp.floataa      )
      a					:      str    =       0
      for batch in tqdm(A    ,		total=A    ,		desc='Evaluating ... '      ):
        a					:      Optional[int]    =       self.data_collator(A      )
        a					:      List[str]    =       self.val_step_fn(A    ,		**A      )
        running_loss += jax_utils.unreplicate(metrics['loss']      )
        i += 1
      return running_loss / i





    def     lowerCamelCase__				(							self     : int    ,		A     : int    ,		A     : Optional[Any]      ):


      '''simple docstring'''




      a					:      List[Any]    =       jax_utils.unreplicate(A      )
      print(F'''SAVING CHECKPOINT IN {save_dir}'''    ,		end=' ... '      )
      self.model_save_fn(A    ,		params=state.params      )
      with open(os.path.join(A    ,		'opt_state.msgpack'      )    ,		'wb'      ) as f:
        f.write(to_bytes(state.opt_state      )      )
      joblib.dump(self.args    ,		os.path.join(A    ,		'args.joblib'      )      )
      joblib.dump(self.data_collator    ,		os.path.join(A    ,		'data_collator.joblib'      )      )
      with open(os.path.join(A    ,		'training_state.json'      )    ,		'w'      ) as f:
        json.dump({'step': state.step.item()}    ,		A      )
      print('DONE'      )




def       snake_case  (A_				:List[Any]     ,      A_				:Optional[int]				):







  '''simple docstring'''

  print(f'''RESTORING CHECKPOINT FROM {save_dir}'''     ,      end=' ... '				)
  with open(os.path.join(A_     ,      'flax_model.msgpack'				)     ,      'rb'				) as f:
    a					:      Union[str, Any]    =       from_bytes(state.params     ,      f.read()				)

  with open(os.path.join(A_     ,      'opt_state.msgpack'				)     ,      'rb'				) as f:
    a					:      List[Any]    =       from_bytes(state.opt_state     ,      f.read()				)

  a					:      Optional[int]    =       joblib.load(os.path.join(A_     ,      'args.joblib'				)				)
  a					:      List[Any]    =       joblib.load(os.path.join(A_     ,      'data_collator.joblib'				)				)

  with open(os.path.join(A_     ,      'training_state.json'				)     ,      'r'				) as f:
    a					:      List[Any]    =       json.load(A_				)
  a					:      int    =       training_state['step']

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




def       snake_case  (A_				:Union[str, Any]     ,      A_				:Optional[int]     ,      A_				:int     ,      A_				:Union[str, Any]				):







  '''simple docstring'''

  a					:      Any    =       num_train_steps - warmup_steps
  a					:      Optional[Any]    =       optax.linear_schedule(init_value=A_     ,      end_value=A_     ,      transition_steps=A_				)
  a					:      List[Any]    =       optax.linear_schedule(init_value=A_     ,      end_value=1E-7     ,      transition_steps=A_				)
  a					:      Optional[int]    =       optax.join_schedules(schedules=[warmup_fn, decay_fn]     ,      boundaries=[warmup_steps]				)
  return lr






def       snake_case  (A_				:str     ,      A_				:int     ,      A_				:List[str]     ,      A_				:Union[str, Any]     ,      A_				:str				):







  '''simple docstring'''

  def weight_decay_mask(A_				:Optional[Any]				):
    a					:      Union[str, Any]    =       traverse_util.flatten_dict(A_				)
    a					:      Optional[int]    =       {k: (v[-1] != 'bias' and v[-2:] != ('LayerNorm', 'scale')) for k, v in params.items()}
    return traverse_util.unflatten_dict(A_				)

  a					:      Any    =       scheduler_fn(A_     ,      A_     ,      A_     ,      A_				)

  a					:      int    =       optax.adamw(learning_rate=A_     ,      weight_decay=A_     ,      mask=A_				)
  return tx, lr












"""simple docstring"""


import unittest

from transformers import SPIECE_UNDERLINE, XLNetTokenizer, XLNetTokenizerFast
from transformers.testing_utils import get_tests_dir, require_sentencepiece, require_tokenizers, slow

from ...test_tokenization_common import TokenizerTesterMixin


_UpperCamelCase    :   List[str]   			=						get_tests_dir('fixtures/test_sentencepiece.model')






@require_sentencepiece
@require_tokenizers
class 						snake_case  (  UpperCAmelCase       ,			unittest.TestCase	):
    __magic_name__   					=					XLNetTokenizer
    __magic_name__   					=					XLNetTokenizerFast
    __magic_name__   					=					True
    __magic_name__   					=					True





    def     lowerCamelCase__				(							self     : List[Any]      ):


      '''simple docstring'''




      super().setUp()

      # We have a SentencePiece fixture for testing
      a					:      Tuple    =       XLNetTokenizer(A    ,		keep_accents=A      )
      tokenizer.sanitize_special_tokens()
      tokenizer.save_pretrained(self.tmpdirname      )





    def     lowerCamelCase__				(							self     : Optional[int]      ):


      '''simple docstring'''




      a					:      Optional[int]    =       '<s>'
      a					:      int    =       1

      self.assertEqual(self.get_tokenizer()._convert_token_to_id(A      )    ,		A      )
      self.assertEqual(self.get_tokenizer()._convert_id_to_token(A      )    ,		A      )





    def     lowerCamelCase__				(							self     : Tuple      ):


      '''simple docstring'''




      a					:      str    =       list(self.get_tokenizer().get_vocab().keys()      )

      self.assertEqual(vocab_keys[0]    ,		'<unk>'      )
      self.assertEqual(vocab_keys[1]    ,		'<s>'      )
      self.assertEqual(vocab_keys[-1]    ,		'<eod>'      )
      self.assertEqual(len(A      )    ,		1_0_0_6      )





    def     lowerCamelCase__				(							self     : Any      ):


      '''simple docstring'''




      self.assertEqual(self.get_tokenizer().vocab_size    ,		1_0_0_0      )





    def     lowerCamelCase__				(							self     : Optional[int]      ):


      '''simple docstring'''




      a					:      Optional[Any]    =       XLNetTokenizer(A    ,		keep_accents=A      )

      a					:      List[str]    =       tokenizer.tokenize('This is a test'      )
      self.assertListEqual(A    ,		['▁This', '▁is', '▁a', '▁t', 'est']      )

      self.assertListEqual(tokenizer.convert_tokens_to_ids(A      )    ,		[2_8_5, 4_6, 1_0, 1_7_0, 3_8_2]      )

      a					:      List[str]    =       tokenizer.tokenize('I was born in 92000, and this is falsé.'      )
      self.assertListEqual(
          A    ,		[
              SPIECE_UNDERLINE + 'I',
              SPIECE_UNDERLINE + 'was',
              SPIECE_UNDERLINE + 'b',
              'or',
              'n',
              SPIECE_UNDERLINE + 'in',
              SPIECE_UNDERLINE + '',
              '9',
              '2',
              '0',
              '0',
              '0',
              ',',
              SPIECE_UNDERLINE + 'and',
              SPIECE_UNDERLINE + 'this',
              SPIECE_UNDERLINE + 'is',
              SPIECE_UNDERLINE + 'f',
              'al',
              's',
              'é',
              '.',
          ]    ,		)
      a					:      Union[str, Any]    =       tokenizer.convert_tokens_to_ids(A      )
      self.assertListEqual(A    ,		[8, 2_1, 8_4, 5_5, 2_4, 1_9, 7, 0, 6_0_2, 3_4_7, 3_4_7, 3_4_7, 3, 1_2, 6_6, 4_6, 7_2, 8_0, 6, 0, 4]      )

      a					:      Tuple    =       tokenizer.convert_ids_to_tokens(A      )
      self.assertListEqual(
          A    ,		[
              SPIECE_UNDERLINE + 'I',
              SPIECE_UNDERLINE + 'was',
              SPIECE_UNDERLINE + 'b',
              'or',
              'n',
              SPIECE_UNDERLINE + 'in',
              SPIECE_UNDERLINE + '',
              '<unk>',
              '2',
              '0',
              '0',
              '0',
              ',',
              SPIECE_UNDERLINE + 'and',
              SPIECE_UNDERLINE + 'this',
              SPIECE_UNDERLINE + 'is',
              SPIECE_UNDERLINE + 'f',
              'al',
              's',
              '<unk>',
              '.',
          ]    ,		)





    def     lowerCamelCase__				(							self     : Union[str, Any]      ):


      '''simple docstring'''




      a					:      Optional[int]    =       XLNetTokenizer(A    ,		do_lower_case=A      )
      a					:      Any    =       tokenizer.tokenize('I was born in 92000, and this is falsé.'      )
      self.assertListEqual(
          A    ,		[
              SPIECE_UNDERLINE + '',
              'i',
              SPIECE_UNDERLINE + 'was',
              SPIECE_UNDERLINE + 'b',
              'or',
              'n',
              SPIECE_UNDERLINE + 'in',
              SPIECE_UNDERLINE + '',
              '9',
              '2',
              '0',
              '0',
              '0',
              ',',
              SPIECE_UNDERLINE + 'and',
              SPIECE_UNDERLINE + 'this',
              SPIECE_UNDERLINE + 'is',
              SPIECE_UNDERLINE + 'f',
              'al',
              'se',
              '.',
          ]    ,		)
      self.assertListEqual(tokenizer.tokenize('H\u00E9llo'      )    ,		['▁he', 'll', 'o']      )





    def     lowerCamelCase__				(							self     : int      ):


      '''simple docstring'''




      a					:      str    =       XLNetTokenizer(A    ,		do_lower_case=A      )
      a					:      List[Any]    =       tokenizer.tokenize('I was born in 92000, and this is falsé.'      )
      self.assertListEqual(
          A    ,		[
              SPIECE_UNDERLINE + 'I',
              SPIECE_UNDERLINE + 'was',
              SPIECE_UNDERLINE + 'b',
              'or',
              'n',
              SPIECE_UNDERLINE + 'in',
              SPIECE_UNDERLINE + '',
              '9',
              '2',
              '0',
              '0',
              '0',
              ',',
              SPIECE_UNDERLINE + 'and',
              SPIECE_UNDERLINE + 'this',
              SPIECE_UNDERLINE + 'is',
              SPIECE_UNDERLINE + 'f',
              'al',
              'se',
              '.',
          ]    ,		)





    @slow
    def     lowerCamelCase__				(							self     : Optional[int]      ):


      '''simple docstring'''




      a					:      Dict    =       XLNetTokenizer.from_pretrained('xlnet-base-cased'      )

      a					:      int    =       tokenizer.encode('sequence builders'    ,		add_special_tokens=A      )
      a					:      List[Any]    =       tokenizer.encode('multi-sequence build'    ,		add_special_tokens=A      )

      a					:      Optional[int]    =       tokenizer.build_inputs_with_special_tokens(A      )
      a					:      Optional[Any]    =       tokenizer.build_inputs_with_special_tokens(A    ,		A      )

      assert encoded_sentence == text + [4, 3]
      assert encoded_pair == text + [4] + text_a + [4, 3]





    @slow
    def     lowerCamelCase__				(							self     : Any      ):


      '''simple docstring'''




      a					:      Union[str, Any]    =       {'input_ids': [[1_7, 2_1_4_4_2, 2_7_0, 1_7, 1_0, 1_4_6_4_5, 3_1_8, 3_4, 1_7, 4_5_4_6, 3_1_4_5, 7_8_7, 1_3, 7_7_5_2, 2_2_0_1_8, 2_3, 2_1, 1_7, 4_5_4_6, 3_1_4_5, 7_8_7, 1_3, 3_3_5_2, 1_4_4_3_1, 1_3, 5_5_0_0, 1_1, 1_1_7_6, 5_8_0, 1_3, 1_6_8_1_9, 4_7_9_7, 2_3, 1_7, 1_0, 1_7_1_3_5, 6_5_8, 1_9, 4_5_7, 7_9_3_2, 1_3, 1_8_4, 1_9, 3_1_5_4, 1_7_1_3_5, 6_4_6_8, 1_9, 1_4_0_4, 1_2_2_6_9, 1_9, 4_2_2_9, 5_3_5_6, 1_6_2_6_4, 4_6, 1_9, 1_7, 2_0_5_4_5, 1_0_3_9_5, 9, 9, 9, 1_1, 2_8, 6_4_2_1, 9_5_3_1, 2_0_7_2_9, 1_7, 1_0, 3_5_3, 1_7_0_2_2, 1_1, 2_1, 6_4_2_1, 9_5_3_1, 1_6_9_4_9, 1_7, 1_0, 1_1_5_0_9, 7_5_3, 1_1, 3_3, 9_5, 2_4_2_1, 7_3_8_5, 9_5_6, 1_4_4_3_1, 2_6_2_6, 2_5, 8_4_2, 7_3_8_5, 4_8_3_6, 2_1, 1_4_2_9, 2_2_7_2, 9_8_5_5, 3_1_2_0, 1_6_1, 2_4_7_3_8, 1_9, 1_3_2_0_3, 6_5_8, 2_1_8, 7_8_7, 2_1, 4_3_0, 1_8_4_8_2, 8_4_7, 2_6_3_7, 9, 4, 3], [5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 3_2_2, 2_2_1_7_8, 2_7, 1_0_6_4, 2_2, 9_5_6, 1_3, 1_1_1_0_1, 1_4_2_9, 5_8_5_4, 2_4_3_1_3, 1_8_9_5_3, 4_0, 4_2_2, 2_4_3_6_6, 6_8, 1_7_5_8, 3_7, 1_0_4_8_3, 1_4_2_5_7, 3_1, 2_0_7, 2_6_3, 2_1, 2_0_3, 3_7_7_3, 2_5, 7_1, 9_7_3_5, 9, 4, 3], [5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 3_2, 2_0_4_9, 3_4_4_2, 1_7, 1_3_8_9_4, 3_3_8_0, 2_3, 9_5, 1_8, 1_7_6_3_4, 2_2_8_8, 9, 4, 3]], 'token_type_ids': [[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2], [3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2], [3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2]], 'attention_mask': [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]}  # noqa: E501
      # fmt: on

      self.tokenizer_integration_test_util(
          expected_encoding=A    ,		model_name='xlnet-base-cased'    ,		revision='c841166438c31ec7ca9a106dee7bb312b73ae511'    ,		)




