



"""simple docstring"""






import copy
import logging
import random

import joblib
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from tqdm import tqdm

from transformers import AdamW, GPTaLMHeadModel, get_linear_schedule_with_warmup


lowerCamelCase				:     Any								=  logging.getLogger(__name__)


def 							A__ (     UpperCamelCase__    ):


  '''simple docstring'''


  random.seed(UpperCamelCase__    )
  np.random.seed(UpperCamelCase__    )
  torch.manual_seed(UpperCamelCase__    )
  torch.cuda.manual_seed_all(UpperCamelCase__    )


def 							A__ (     UpperCamelCase__       ,							UpperCamelCase__       ,							UpperCamelCase__    ):


  '''simple docstring'''


  model.eval()
  _SCREAMING_SNAKE_CASE        =			next(model.parameters()    ).device
  _SCREAMING_SNAKE_CASE        =			1
  _SCREAMING_SNAKE_CASE        =			torch.zeros((eval_batch_size, context_len)       ,							dtype=torch.long       ,							device=UpperCamelCase__    )
  _SCREAMING_SNAKE_CASE        =			DataLoader(UpperCamelCase__       ,							shuffle=UpperCamelCase__       ,							batch_size=UpperCamelCase__    )
  _SCREAMING_SNAKE_CASE        =			torch.zeros(1       ,							device=UpperCamelCase__    )
  _SCREAMING_SNAKE_CASE        =			0
  for batch in eval_dataloader:
    batch.to(UpperCamelCase__    )
    # pad
    context.zero_()
    for i in range(UpperCamelCase__    ):
      _SCREAMING_SNAKE_CASE        =			batch[i]
    _SCREAMING_SNAKE_CASE        =			model(UpperCamelCase__       ,							labels=UpperCamelCase__    )
    eval_loss += outputs[0].sum().item()
    nb_eval_examples += batch.size(0    )
  _SCREAMING_SNAKE_CASE        =			eval_loss / nb_eval_examples
  _SCREAMING_SNAKE_CASE        =			torch.exp(UpperCamelCase__    )
  model.train()
  return perplexity


def 							A__ (     UpperCamelCase__="gpt2"    ):


  '''simple docstring'''


  _SCREAMING_SNAKE_CASE        =			GPTaLMHeadModel.from_pretrained(UpperCamelCase__       ,							output_hidden_states=UpperCamelCase__    )
  torch.save(model.state_dict()       ,							model_name + '''local.pt'''    )
  return model


def 							A__ (     UpperCamelCase__       ,							UpperCamelCase__       ,							UpperCamelCase__    ):


  '''simple docstring'''


  _SCREAMING_SNAKE_CASE        =			copy.deepcopy(UpperCamelCase__    )
  model.to(UpperCamelCase__    )

  _SCREAMING_SNAKE_CASE        =			['''bias''', '''LayerNorm.weight''']
  _SCREAMING_SNAKE_CASE        =			[
      {
          '''params''': [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay    )],
          '''weight_decay''': 0.0,
      },
      {'''params''': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay    )], '''weight_decay''': 0.0},
  ]
  _SCREAMING_SNAKE_CASE        =			AdamW(UpperCamelCase__       ,							lr=5e-5       ,							eps=1e-8    )
  _SCREAMING_SNAKE_CASE        =			get_linear_schedule_with_warmup(UpperCamelCase__       ,							0       ,							UpperCamelCase__    )
  torch.cuda.empty_cache()
  return model, lm_optimizer, lm_scheduler


def 							A__ (     UpperCamelCase__       ,							UpperCamelCase__       ,							UpperCamelCase__       ,							UpperCamelCase__    ):


  '''simple docstring'''


  _SCREAMING_SNAKE_CASE        =			np.array(UpperCamelCase__    ).mean()
  _SCREAMING_SNAKE_CASE        =			np.array(UpperCamelCase__    ).std()
  _SCREAMING_SNAKE_CASE        =			(real_perps - avg) / std
  _SCREAMING_SNAKE_CASE        =			list(zip(UpperCamelCase__       ,							UpperCamelCase__       ,							UpperCamelCase__    )    )
  joblib.dump(UpperCamelCase__       ,							UpperCamelCase__    )


def 							A__ (     UpperCamelCase__       ,							UpperCamelCase__       ,							UpperCamelCase__       ,							UpperCamelCase__       ,							UpperCamelCase__       ,							UpperCamelCase__       ,							UpperCamelCase__       ,							UpperCamelCase__="dev.jbl"       ,							UpperCamelCase__=recopy_gpta       ,							):


  '''simple docstring'''


  _SCREAMING_SNAKE_CASE        =			[]
  _SCREAMING_SNAKE_CASE        =			[]
  _SCREAMING_SNAKE_CASE        =			[]

  # Initialize the transformer model
  _SCREAMING_SNAKE_CASE        =			copy.deepcopy(UpperCamelCase__    )
  orig_model.to(device='''cpu'''    )
  torch.cuda.empty_cache()

  # Compute perplexity of initial transformer model for comparison
  model.train()
  _SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE        =			recopy_model(UpperCamelCase__       ,							UpperCamelCase__       ,							UpperCamelCase__    )

  for step in tqdm(range(UpperCamelCase__    )    ):
    _SCREAMING_SNAKE_CASE        =			torch.zeros((1, context_len)       ,							dtype=torch.long       ,							device=UpperCamelCase__    )
    _SCREAMING_SNAKE_CASE        =			random.choice(UpperCamelCase__    )
    _SCREAMING_SNAKE_CASE        =			random.randint(0       ,							len(story[0]    ) - context_len - 1    )
    _SCREAMING_SNAKE_CASE        =			story[0][start : start + context_len]
    lm_optimizer.zero_grad()
    _SCREAMING_SNAKE_CASE        =			model(UpperCamelCase__       ,							labels=UpperCamelCase__    )
    _SCREAMING_SNAKE_CASE        =			outputs[0]
    _SCREAMING_SNAKE_CASE        =			compute_perplexity(UpperCamelCase__       ,							UpperCamelCase__       ,							UpperCamelCase__    )
    model.train()
    lm_loss.backward()
    # Do LM backprop
    torch.nn.utils.clip_grad_norm_(model.parameters()       ,							3.0    )
    lm_optimizer.step()
    lm_scheduler.step()  # Update learning rate schedule

    # Compute perplexity after back-propagating on the selected context
    _SCREAMING_SNAKE_CASE        =			compute_perplexity(UpperCamelCase__       ,							UpperCamelCase__       ,							UpperCamelCase__    )

    # Periodically save the stored (X, IG(X)) pairs
    if step % 1_000 == 0 and step > 1:
      intermittent_save(UpperCamelCase__       ,							UpperCamelCase__       ,							UpperCamelCase__       ,							UpperCamelCase__    )

    # Reset the pretrained model to the original pretrained GPT-2 weights after each iteration
    _SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE    ,     _SCREAMING_SNAKE_CASE        =			recopy_model(UpperCamelCase__       ,							UpperCamelCase__       ,							UpperCamelCase__    )

    past_perps.append(past_perp.item()    )
    real_perps.append(orig_perp - real_perp.item()    )
    contexts.append(np.array(context.cpu()    )    )

  intermittent_save(UpperCamelCase__       ,							UpperCamelCase__       ,							UpperCamelCase__       ,							UpperCamelCase__    )


def 							A__ (     UpperCamelCase__       ,							UpperCamelCase__="data/tokenized_stories_train_wikitext103.jbl"       ,							UpperCamelCase__=100       ,							UpperCamelCase__=1_026       ,							UpperCamelCase__=True    ):


  '''simple docstring'''


  _SCREAMING_SNAKE_CASE        =			joblib.load(UpperCamelCase__    )
  print('''data loaded'''    )
  _SCREAMING_SNAKE_CASE        =			[]
  if trim:
    for i, example in enumerate(UpperCamelCase__    ):
      if len(example[0]    ) > min_len:
        _SCREAMING_SNAKE_CASE        =			random.randint(0       ,							len(example[0]    ) - context_len - 1    )
        objective_set.append(example[0, start : start + context_len]    )
      if len(UpperCamelCase__    ) >= number:
        break
    _SCREAMING_SNAKE_CASE        =			[]
    for j in range(i + 1       ,							len(UpperCamelCase__    )    ):
      if len(data[j][0]    ) > min_len:
        train_data.append(data[j]    )
  else:
    _SCREAMING_SNAKE_CASE        =			data[0:number]
    _SCREAMING_SNAKE_CASE        =			data[number:]

  joblib.dump(UpperCamelCase__       ,							'''objective_set.jbl'''    )
  print('''objective set saved'''    )
  return train_data, objective_set


def 							A__ (     UpperCamelCase__       ,							UpperCamelCase__       ,							UpperCamelCase__       ,							UpperCamelCase__       ,							UpperCamelCase__=50       ,							UpperCamelCase__="secondary_learner.pt"    ):


  '''simple docstring'''


  _SCREAMING_SNAKE_CASE        =			torch.device('''cuda:0''' if torch.cuda.is_available() else '''cpu'''    )
  # We will use the first 512 pairs from our dataset as a test set for
  # our secondary learner and the rest to train
  _SCREAMING_SNAKE_CASE        =			train_dataset[:512]
  _SCREAMING_SNAKE_CASE        =			train_dataset[512:]
  _SCREAMING_SNAKE_CASE        =			DataLoader(UpperCamelCase__       ,							shuffle=UpperCamelCase__       ,							batch_size=UpperCamelCase__    )
  _SCREAMING_SNAKE_CASE        =			DataLoader(UpperCamelCase__       ,							shuffle=UpperCamelCase__       ,							batch_size=UpperCamelCase__    )

  # secondary learner model set up
  _SCREAMING_SNAKE_CASE        =			nn.MSELoss()
  _SCREAMING_SNAKE_CASE        =			nn.MSELoss(reduction='''sum'''    )
  secondary_learner.to(UpperCamelCase__    )
  _SCREAMING_SNAKE_CASE        =			torch.optim.Adam(secondary_learner.parameters()       ,							lr=0.0_00_01    )
  secondary_learner.train()

  # TODO in original code this is written as number of actual batches seen
  # not number of items seen but other places it is number of items instead.
  # improve consistency! changed this to epochs for clarity
  _SCREAMING_SNAKE_CASE        =			float('''inf'''    )
  # Iterate through batches until we've used max_steps batches
  for epoch in range(int(UpperCamelCase__    )    ):
    _SCREAMING_SNAKE_CASE        =			0.0
    secondary_learner.train()
    for step, batch in enumerate(UpperCamelCase__    ):
      _SCREAMING_SNAKE_CASE        =			batch[0].to(UpperCamelCase__    )
      _SCREAMING_SNAKE_CASE        =			batch[1].to(UpperCamelCase__    )
      _SCREAMING_SNAKE_CASE        =			secondary_learner(UpperCamelCase__    )
      q_optimizer.zero_grad()
      _SCREAMING_SNAKE_CASE        =			loss(UpperCamelCase__       ,							real_q.float()    )
      q_loss.backward()
      q_optimizer.step()
      tr_q_loss += q_loss.item()

      # model trains fairly quickly so we won't wait for a full epoch
      # eval is triggered at eval_freq and end of epochs
      if (step % eval_freq == 0 and step > 0) or ((step + 1) == len(UpperCamelCase__    )):
        _SCREAMING_SNAKE_CASE        =			tr_q_loss / (step + 1)

        secondary_learner.eval()
        _SCREAMING_SNAKE_CASE        =			0.0
        _SCREAMING_SNAKE_CASE        =			0.0
        _SCREAMING_SNAKE_CASE        =			[]
        _SCREAMING_SNAKE_CASE        =			[]
        # Compute performance of the secondary learner after this batch
        for stepa, batcha in enumerate(UpperCamelCase__    ):
          _SCREAMING_SNAKE_CASE        =			batcha[0].to(UpperCamelCase__    )
          _SCREAMING_SNAKE_CASE        =			batcha[1].to(UpperCamelCase__    )
          _SCREAMING_SNAKE_CASE        =			secondary_learner(UpperCamelCase__    )
          q_lossa += test_loss(UpperCamelCase__       ,							UpperCamelCase__    ).item()
          sum_qa += torch.sum(UpperCamelCase__    ).item()
          for ei, i in enumerate(predicted_qa.cpu().detach().numpy()    ):
            predicted.append(i.item()    )
          for ei, i in enumerate(real_qa.cpu().detach().numpy()    ):
            actual.append(i.item()    )

        q_lossa /= len(UpperCamelCase__    )
        print(
            '''Epoch: '''       ,							UpperCamelCase__       ,							'''step: '''       ,							UpperCamelCase__       ,							'''Avg. q:'''       ,							sum_qa / len(UpperCamelCase__    )       ,							'''Train Loss: '''       ,							UpperCamelCase__       ,							'''Test Loss: '''       ,							UpperCamelCase__       ,							)
        if q_lossa < best_test_loss:
          joblib.dump((predicted, actual)       ,							'''pred_vs_actual.jbl'''    )
          torch.save(secondary_learner.state_dict()       ,							UpperCamelCase__    )
          _SCREAMING_SNAKE_CASE        =			q_lossa

      secondary_learner.train()
  return secondary_learner







class    __snake_case(					nn.Module       ):




    def __init__(   self    ,							A_			):





      '''simple docstring'''



      super(a__    ,							self			).__init__()
      _SCREAMING_SNAKE_CASE        =			model.transformer.wte
      _SCREAMING_SNAKE_CASE        =			copy.deepcopy(model.transformer.wte.weight			)
      _SCREAMING_SNAKE_CASE        =			nn.Convad(self.embeddings.weight.size(1			)    ,							256    ,							3    ,							padding=1			)
      _SCREAMING_SNAKE_CASE        =			nn.Sequential(nn.Linear(256    ,							32			)    ,							nn.Dropout(p=0.1			)    ,							nn.Linear(32    ,							32			)    ,							nn.Linear(32    ,							1			)			)




    def   A   (   self    ,							A_			):





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE        =			torch.max(self.conv(self.embeddings(a__			).squeeze(1			).transpose(1    ,							2			)			)    ,							2			)[0]
      _SCREAMING_SNAKE_CASE        =			self.fc(a__			)
      return qs.squeeze(1			)




    @classmethod
    def   A   (   cls    ,							A_    ,							A_			):





      '''simple docstring'''



      _SCREAMING_SNAKE_CASE        =			cls(a__			)  # this calls __init__
      _SCREAMING_SNAKE_CASE        =			torch.load(a__			)
      secondary_learner.load_state_dict(a__			)
      _SCREAMING_SNAKE_CASE        =			model.transformer.wte
      _SCREAMING_SNAKE_CASE        =			copy.deepcopy(model.transformer.wte.weight			)
      return secondary_learner







"""simple docstring"""






from __future__ import annotations

import unittest

from transformers import is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
    import tensorflow as tf

    from transformers import (
        TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFXLMForMultipleChoice,
        TFXLMForQuestionAnsweringSimple,
        TFXLMForSequenceClassification,
        TFXLMForTokenClassification,
        TFXLMModel,
        TFXLMWithLMHeadModel,
        XLMConfig,
    )



class 						lowerCAmelCase       :






       def __init__( self     ,	a__     ,	):
              _UpperCAmelCase			     =  parent
              _UpperCAmelCase			     =  13
              _UpperCAmelCase			     =  7
              _UpperCAmelCase			     =  True
              _UpperCAmelCase			     =  True
              _UpperCAmelCase			     =  True
              _UpperCAmelCase			     =  True
              _UpperCAmelCase			     =  True
              _UpperCAmelCase			     =  False
              _UpperCAmelCase			     =  False
              _UpperCAmelCase			     =  False
              _UpperCAmelCase			     =  2
              _UpperCAmelCase			     =  99
              _UpperCAmelCase			     =  0
              _UpperCAmelCase			     =  32
              _UpperCAmelCase			     =  2
              _UpperCAmelCase			     =  4
              _UpperCAmelCase			     =  0.1
              _UpperCAmelCase			     =  0.1
              _UpperCAmelCase			     =  5_12
              _UpperCAmelCase			     =  16
              _UpperCAmelCase			     =  2
              _UpperCAmelCase			     =  0.02
              _UpperCAmelCase			     =  3
              _UpperCAmelCase			     =  4
              _UpperCAmelCase			     =  'last'
              _UpperCAmelCase			     =  True
              _UpperCAmelCase			     =  None
              _UpperCAmelCase			     =  0






       def 				__A			( self  ):
              _UpperCAmelCase			     =  ids_tensor([self.batch_size, self.seq_length]     ,	self.vocab_size  )
              _UpperCAmelCase			     =  random_attention_mask([self.batch_size, self.seq_length]     ,	dtype=tf.floataa  )

              _UpperCAmelCase			     =  None
              if self.use_input_lengths:
                     _UpperCAmelCase			     =  (
                         ids_tensor([self.batch_size]     ,	vocab_size=2  ) + self.seq_length - 2
                     )  # small variation of seq_length

              _UpperCAmelCase			     =  None
              if self.use_token_type_ids:
                     _UpperCAmelCase			     =  ids_tensor([self.batch_size, self.seq_length]     ,	self.n_langs  )

              _UpperCAmelCase			     =  None
              _UpperCAmelCase			     =  None
              _UpperCAmelCase			     =  None
              if self.use_labels:
                     _UpperCAmelCase			     =  ids_tensor([self.batch_size]     ,	self.type_sequence_label_size  )
                     _UpperCAmelCase			     =  ids_tensor([self.batch_size, self.seq_length]     ,	self.num_labels  )
                     _UpperCAmelCase			     =  ids_tensor([self.batch_size]     ,	2     ,	dtype=tf.floataa  )
                     _UpperCAmelCase			     =  ids_tensor([self.batch_size]     ,	self.num_choices  )

              _UpperCAmelCase			     =  XLMConfig(
                  vocab_size=self.vocab_size     ,	n_special=self.n_special     ,	emb_dim=self.hidden_size     ,	n_layers=self.num_hidden_layers     ,	n_heads=self.num_attention_heads     ,	dropout=self.hidden_dropout_prob     ,	attention_dropout=self.attention_probs_dropout_prob     ,	gelu_activation=self.gelu_activation     ,	sinusoidal_embeddings=self.sinusoidal_embeddings     ,	asm=self.asm     ,	causal=self.causal     ,	n_langs=self.n_langs     ,	max_position_embeddings=self.max_position_embeddings     ,	initializer_range=self.initializer_range     ,	summary_type=self.summary_type     ,	use_proj=self.use_proj     ,	bos_token_id=self.bos_token_id     ,	)

              return (
                  config,
                  input_ids,
                  token_type_ids,
                  input_lengths,
                  sequence_labels,
                  token_labels,
                  is_impossible_labels,
                  choice_labels,
                  input_mask,
              )






       def 				__A			( self     ,	a__     ,	a__     ,	a__     ,	a__     ,	a__     ,	a__     ,	a__     ,	a__     ,	a__     ,	):
              _UpperCAmelCase			     =  TFXLMModel(config=a__  )
              _UpperCAmelCase			     =  {'input_ids': input_ids, 'lengths': input_lengths, 'langs': token_type_ids}
              _UpperCAmelCase			     =  model(a__  )

              _UpperCAmelCase			     =  [input_ids, input_mask]
              _UpperCAmelCase			     =  model(a__  )
              self.parent.assertEqual(result.last_hidden_state.shape     ,	(self.batch_size, self.seq_length, self.hidden_size)  )






       def 				__A			( self     ,	a__     ,	a__     ,	a__     ,	a__     ,	a__     ,	a__     ,	a__     ,	a__     ,	a__     ,	):
              _UpperCAmelCase			     =  TFXLMWithLMHeadModel(a__  )

              _UpperCAmelCase			     =  {'input_ids': input_ids, 'lengths': input_lengths, 'langs': token_type_ids}
              _UpperCAmelCase			     =  model(a__  )

              _UpperCAmelCase			     =  outputs

              self.parent.assertEqual(result.logits.shape     ,	(self.batch_size, self.seq_length, self.vocab_size)  )






       def 				__A			( self     ,	a__     ,	a__     ,	a__     ,	a__     ,	a__     ,	a__     ,	a__     ,	a__     ,	a__     ,	):
              _UpperCAmelCase			     =  TFXLMForQuestionAnsweringSimple(a__  )

              _UpperCAmelCase			     =  {'input_ids': input_ids, 'lengths': input_lengths}

              _UpperCAmelCase			     =  model(a__  )

              self.parent.assertEqual(result.start_logits.shape     ,	(self.batch_size, self.seq_length)  )
              self.parent.assertEqual(result.end_logits.shape     ,	(self.batch_size, self.seq_length)  )






       def 				__A			( self     ,	a__     ,	a__     ,	a__     ,	a__     ,	a__     ,	a__     ,	a__     ,	a__     ,	a__     ,	):
              _UpperCAmelCase			     =  TFXLMForSequenceClassification(a__  )

              _UpperCAmelCase			     =  {'input_ids': input_ids, 'lengths': input_lengths}

              _UpperCAmelCase			     =  model(a__  )

              self.parent.assertEqual(result.logits.shape     ,	(self.batch_size, self.type_sequence_label_size)  )






       def 				__A			( self     ,	a__     ,	a__     ,	a__     ,	a__     ,	a__     ,	a__     ,	a__     ,	a__     ,	a__     ,	):
              _UpperCAmelCase			     =  self.num_labels
              _UpperCAmelCase			     =  TFXLMForTokenClassification(config=a__  )
              _UpperCAmelCase			     =  {'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
              _UpperCAmelCase			     =  model(a__  )
              self.parent.assertEqual(result.logits.shape     ,	(self.batch_size, self.seq_length, self.num_labels)  )






       def 				__A			( self     ,	a__     ,	a__     ,	a__     ,	a__     ,	a__     ,	a__     ,	a__     ,	a__     ,	a__     ,	):
              _UpperCAmelCase			     =  self.num_choices
              _UpperCAmelCase			     =  TFXLMForMultipleChoice(config=a__  )
              _UpperCAmelCase			     =  tf.tile(tf.expand_dims(a__     ,	1  )     ,	(1, self.num_choices, 1)  )
              _UpperCAmelCase			     =  tf.tile(tf.expand_dims(a__     ,	1  )     ,	(1, self.num_choices, 1)  )
              _UpperCAmelCase			     =  tf.tile(tf.expand_dims(a__     ,	1  )     ,	(1, self.num_choices, 1)  )
              _UpperCAmelCase			     =  {
                  'input_ids': multiple_choice_inputs_ids,
                  'attention_mask': multiple_choice_input_mask,
                  'token_type_ids': multiple_choice_token_type_ids,
              }
              _UpperCAmelCase			     =  model(a__  )
              self.parent.assertEqual(result.logits.shape     ,	(self.batch_size, self.num_choices)  )






       def 				__A			( self  ):
              _UpperCAmelCase			     =  self.prepare_config_and_inputs()
              (
                  (
                  _UpperCAmelCase
              )     ,				(
                  _UpperCAmelCase
              )     ,				(
                  _UpperCAmelCase
              )     ,				(
                  _UpperCAmelCase
              )     ,				(
                  _UpperCAmelCase
              )     ,				(
                  _UpperCAmelCase
              )     ,				(
                  _UpperCAmelCase
              )     ,				(
                  _UpperCAmelCase
              )     ,				(
                  _UpperCAmelCase
              )     ,				
              )			     =  config_and_inputs
              _UpperCAmelCase			     =  {
                  'input_ids': input_ids,
                  'token_type_ids': token_type_ids,
                  'langs': token_type_ids,
                  'lengths': input_lengths,
              }
              return config, inputs_dict



@require_tf
class 						lowerCAmelCase       (   snake_case      ,						snake_case      ,						unittest.TestCase				):
       lowerCAmelCase__    			=	(
           (
               TFXLMModel,
               TFXLMWithLMHeadModel,
               TFXLMForSequenceClassification,
               TFXLMForQuestionAnsweringSimple,
               TFXLMForTokenClassification,
               TFXLMForMultipleChoice,
           )
           if is_tf_available()
           else ()
       )
       lowerCAmelCase__    			=	(
           (TFXLMWithLMHeadModel,) if is_tf_available() else ()
       )  # TODO (PVP): Check other models whether language generation is also applicable
       lowerCAmelCase__    			=	(
           {
               """feature-extraction""": TFXLMModel,
               """fill-mask""": TFXLMWithLMHeadModel,
               """question-answering""": TFXLMForQuestionAnsweringSimple,
               """text-classification""": TFXLMForSequenceClassification,
               """text-generation""": TFXLMWithLMHeadModel,
               """token-classification""": TFXLMForTokenClassification,
               """zero-shot""": TFXLMForSequenceClassification,
           }
           if is_tf_available()
           else {}
       )
       lowerCAmelCase__    			=	False
       lowerCAmelCase__    			=	False






       def 				__A			( self     ,	a__     ,	a__     ,	a__     ,	a__     ,	a__  ):
              if (
                  pipeline_test_casse_name == "QAPipelineTests"
                  and tokenizer_name is not None
                  and not tokenizer_name.endswith('Fast'  )
              ):
                     # `QAPipelineTests` fails for a few models when the slower tokenizer are used.
                     # (The slower tokenizers were never used for pipeline tests before the pipeline testing rework)
                     # TODO: check (and possibly fix) the `QAPipelineTests` with slower tokenizer
                     return True

              return False






       def 				__A			( self  ):
              _UpperCAmelCase			     =  TFXLMModelTester(self  )
              _UpperCAmelCase			     =  ConfigTester(self     ,	config_class=a__     ,	emb_dim=37  )






       def 				__A			( self  ):
              self.config_tester.run_common_tests()






       def 				__A			( self  ):
              _UpperCAmelCase			     =  self.model_tester.prepare_config_and_inputs()
              self.model_tester.create_and_check_xlm_model(*a__  )






       def 				__A			( self  ):
              _UpperCAmelCase			     =  self.model_tester.prepare_config_and_inputs()
              self.model_tester.create_and_check_xlm_lm_head(*a__  )






       def 				__A			( self  ):
              _UpperCAmelCase			     =  self.model_tester.prepare_config_and_inputs()
              self.model_tester.create_and_check_xlm_qa(*a__  )






       def 				__A			( self  ):
              _UpperCAmelCase			     =  self.model_tester.prepare_config_and_inputs()
              self.model_tester.create_and_check_xlm_sequence_classif(*a__  )






       def 				__A			( self  ):
              _UpperCAmelCase			     =  self.model_tester.prepare_config_and_inputs()
              self.model_tester.create_and_check_xlm_for_token_classification(*a__  )






       def 				__A			( self  ):
              _UpperCAmelCase			     =  self.model_tester.prepare_config_and_inputs()
              self.model_tester.create_and_check_xlm_for_multiple_choice(*a__  )






       @slow
       def 				__A			( self  ):
              for model_name in TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                     _UpperCAmelCase			     =  TFXLMModel.from_pretrained(a__  )
                     self.assertIsNotNone(a__  )



@require_tf
class 						lowerCAmelCase       (   unittest.TestCase				):






       @slow
       def 				__A			( self  ):
              _UpperCAmelCase			     =  TFXLMWithLMHeadModel.from_pretrained('xlm-mlm-en-2048'  )
              _UpperCAmelCase			     =  tf.convert_to_tensor([[14, 4_47]]     ,	dtype=tf.intaa  )  # the president
              _UpperCAmelCase			     =  [
                  14,
                  4_47,
                  14,
                  4_47,
                  14,
                  4_47,
                  14,
                  4_47,
                  14,
                  4_47,
                  14,
                  4_47,
                  14,
                  4_47,
                  14,
                  4_47,
                  14,
                  4_47,
                  14,
                  4_47,
              ]  # the president the president the president the president the president the president the president the president the president the president
              # TODO(PVP): this and other input_ids I tried for generation give pretty bad results. Not sure why. Model might just not be made for auto-regressive inference
              _UpperCAmelCase			     =  model.generate(a__     ,	do_sample=a__  )
              self.assertListEqual(output_ids[0].numpy().tolist()     ,	a__  )






