

'''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 		__SCREAMING_SNAKE_CASE       ( lowercase__			):
				lowerCamelCase_							=						42
				lowerCamelCase_							=						jnp.floataa
				lowerCamelCase_							=						True







				def    lowerCamelCase_  (					self    :		Tuple      ):
										'''simple docstring'''
										super().setup()
										lowercase   :       Tuple								=nn.Dense(5		,				dtype=self.dtype      )







				def __call__(					self    :		Optional[Any]		,				*UpperCAmelCase__    :		str		,				**UpperCAmelCase__    :		Any      ):
										'''simple docstring'''
										lowercase   :       Optional[int]								=super().__call__(*UpperCAmelCase__		,				**UpperCAmelCase__      )
										lowercase   :       Any								=self.cls(outputs[2]      )
										return outputs[:2] + (cls_out,)





class 		__SCREAMING_SNAKE_CASE       ( lowercase__			):
				lowerCamelCase_							=						FlaxBigBirdForNaturalQuestionsModule






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

						lowercase   :       Optional[int]								=partial(__magic_name__      ,   reduction=jnp.mean	)
						lowercase   :       Union[str, Any]								=cross_entropy(__magic_name__      ,   __magic_name__	)
						lowercase   :       List[Any]								=cross_entropy(__magic_name__      ,   __magic_name__	)
						lowercase   :       Tuple								=cross_entropy(__magic_name__      ,   __magic_name__	)
						return (start_loss + end_loss + pooled_loss) / 3





@dataclass
class 		__SCREAMING_SNAKE_CASE       :
				lowerCamelCase_							=						"google/bigbird-roberta-base"
				lowerCamelCase_							=						30_00
				lowerCamelCase_							=						1_05_00

				lowerCamelCase_							=						1_28
				lowerCamelCase_							=						3

				lowerCamelCase_							=						1
				lowerCamelCase_							=						5

				# tx_args
				lowerCamelCase_							=						3E-5
				lowerCamelCase_							=						0.0
				lowerCamelCase_							=						2_00_00
				lowerCamelCase_							=						0.0_0_9_5

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







				def    lowerCamelCase_  (					self    :		Union[str, Any]      ):
										'''simple docstring'''
										os.makedirs(self.base_dir		,				exist_ok=UpperCAmelCase__      )
										lowercase   :       Tuple								=os.path.join(self.base_dir		,				self.save_dir      )
										lowercase   :       List[str]								=self.batch_size_per_device * jax.device_count()





@dataclass
class 		__SCREAMING_SNAKE_CASE       :
				lowerCamelCase_							=						42
				lowerCamelCase_							=						40_96  # no dynamic padding on TPUs







				def __call__(					self    :		List[str]		,				UpperCAmelCase__    :		Union[str, Any]      ):
										'''simple docstring'''
										lowercase   :       List[str]								=self.collate_fn(UpperCAmelCase__      )
										lowercase   :       Tuple								=jax.tree_util.tree_map(UpperCAmelCase__		,				UpperCAmelCase__      )
										return batch







				def    lowerCamelCase_  (					self    :		Optional[int]		,				UpperCAmelCase__    :		Union[str, Any]      ):
										'''simple docstring'''
										lowercase		,  lowercase   :       str								=self.fetch_inputs(features['''input_ids''']      )
										lowercase   :       Optional[int]								={
										    '''input_ids''': jnp.array(UpperCAmelCase__		,				dtype=jnp.intaa      ),
										    '''attention_mask''': jnp.array(UpperCAmelCase__		,				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    :		Dict		,				UpperCAmelCase__    :		list      ):
										'''simple docstring'''
										lowercase   :       int								=[self._fetch_inputs(UpperCAmelCase__      ) for ids in input_ids]
										return zip(*UpperCAmelCase__      )







				def    lowerCamelCase_  (					self    :		Dict		,				UpperCAmelCase__    :		list      ):
										'''simple docstring'''
										lowercase   :       Tuple								=[1 for _ in range(len(UpperCAmelCase__      )      )]
										while len(UpperCAmelCase__      ) < self.max_length:
																input_ids.append(self.pad_id      )
																attention_mask.append(0      )
										return input_ids, attention_mask






def 							_lowerCAmelCase (      __magic_name__	:    Union[str, Any]      ,   __magic_name__	:    List[str]      ,   __magic_name__	:    str=None	)							->     List[str]:
						if seed is not None:
												lowercase   :       Optional[int]								=dataset.shuffle(seed=__magic_name__	)
						for i in range(len(__magic_name__	) // batch_size	):
												lowercase   :       str								=dataset[i * batch_size : (i + 1) * batch_size]
												yield dict(__magic_name__	)






@partial(jax.pmap      ,   axis_name='''batch'''	)
def 							_lowerCAmelCase (      __magic_name__	:    Tuple      ,   __magic_name__	:    int      ,   **__magic_name__	:    str	)							->     Optional[int]:
						def loss_fn(__magic_name__	:    str	):
												lowercase   :       str								=model_inputs.pop('''start_labels'''	)
												lowercase   :       Any								=model_inputs.pop('''end_labels'''	)
												lowercase   :       Dict								=model_inputs.pop('''pooled_labels'''	)

												lowercase   :       Dict								=state.apply_fn(**__magic_name__      ,   params=__magic_name__      ,   dropout_rng=__magic_name__      ,   train=__magic_name__	)
												lowercase		,  lowercase		,  lowercase   :       List[Any]								=outputs

												return state.loss_fn(
												    __magic_name__      ,   __magic_name__      ,   __magic_name__      ,   __magic_name__      ,   __magic_name__      ,   __magic_name__      ,   )

						lowercase		,  lowercase   :       List[Any]								=jax.random.split(__magic_name__	)
						lowercase   :       Union[str, Any]								=jax.value_and_grad(__magic_name__	)
						lowercase		,  lowercase   :       Tuple								=grad_fn(state.params	)
						lowercase   :       str								=jax.lax.pmean({'''loss''': loss}      ,   axis_name='''batch'''	)
						lowercase   :       List[Any]								=jax.lax.pmean(__magic_name__      ,   '''batch'''	)

						lowercase   :       str								=state.apply_gradients(grads=__magic_name__	)
						return state, metrics, new_drp_rng






@partial(jax.pmap      ,   axis_name='''batch'''	)
def 							_lowerCAmelCase (      __magic_name__	:    str      ,   **__magic_name__	:    Optional[Any]	)							->     int:
						lowercase   :       int								=model_inputs.pop('''start_labels'''	)
						lowercase   :       Union[str, Any]								=model_inputs.pop('''end_labels'''	)
						lowercase   :       int								=model_inputs.pop('''pooled_labels'''	)

						lowercase   :       Dict								=state.apply_fn(**__magic_name__      ,   params=state.params      ,   train=__magic_name__	)
						lowercase		,  lowercase		,  lowercase   :       Any								=outputs

						lowercase   :       Optional[Any]								=state.loss_fn(__magic_name__      ,   __magic_name__      ,   __magic_name__      ,   __magic_name__      ,   __magic_name__      ,   __magic_name__	)
						lowercase   :       Optional[Any]								=jax.lax.pmean({'''loss''': loss}      ,   axis_name='''batch'''	)
						return metrics





class 		__SCREAMING_SNAKE_CASE       ( train_state.TrainState			):
				lowerCamelCase_							=						struct.field(pytree_node=lowercase__			)





@dataclass
class 		__SCREAMING_SNAKE_CASE       :
				lowerCamelCase_							=						42
				lowerCamelCase_							=						42
				lowerCamelCase_							=						42
				lowerCamelCase_							=						42
				lowerCamelCase_							=						42
				lowerCamelCase_							=						42
				lowerCamelCase_							=						None







				def    lowerCamelCase_  (					self    :		Optional[int]		,				UpperCAmelCase__    :		List[Any]		,				UpperCAmelCase__    :		int		,				UpperCAmelCase__    :		str		,				UpperCAmelCase__    :		List[str]=None      ):
										'''simple docstring'''
										lowercase   :       Union[str, Any]								=model.params
										lowercase   :       Optional[int]								=TrainState.create(
										    apply_fn=model.__call__		,				params=UpperCAmelCase__		,				tx=UpperCAmelCase__		,				loss_fn=UpperCAmelCase__		,				)
										if ckpt_dir is not None:
																lowercase		,  lowercase		,  lowercase		,  lowercase		,  lowercase   :       List[str]								=restore_checkpoint(UpperCAmelCase__		,				UpperCAmelCase__      )
																lowercase   :       Optional[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,
																}
																lowercase		,  lowercase   :       Optional[int]								=build_tx(**UpperCAmelCase__      )
																lowercase   :       str								=train_state.TrainState(
																    step=UpperCAmelCase__		,				apply_fn=model.__call__		,				params=UpperCAmelCase__		,				tx=UpperCAmelCase__		,				opt_state=UpperCAmelCase__		,				)
																lowercase   :       int								=args
																lowercase   :       Dict								=data_collator
																lowercase   :       Optional[int]								=lr
																lowercase   :       int								=params
										lowercase   :       Optional[Any]								=jax_utils.replicate(UpperCAmelCase__      )
										return state







				def    lowerCamelCase_  (					self    :		List[Any]		,				UpperCAmelCase__    :		Dict		,				UpperCAmelCase__    :		Union[str, Any]		,				UpperCAmelCase__    :		Optional[int]      ):
										'''simple docstring'''
										lowercase   :       List[Any]								=self.args
										lowercase   :       Dict								=len(UpperCAmelCase__      ) // args.batch_size

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

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

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







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







				def    lowerCamelCase_  (					self    :		Tuple		,				UpperCAmelCase__    :		Dict		,				UpperCAmelCase__    :		str      ):
										'''simple docstring'''
										lowercase   :       Tuple								=jax_utils.unreplicate(UpperCAmelCase__      )
										print(F'''SAVING CHECKPOINT IN {save_dir}'''		,				end=''' ... '''      )
										self.model_save_fn(UpperCAmelCase__		,				params=state.params      )
										with open(os.path.join(UpperCAmelCase__		,				'''opt_state.msgpack'''      )		,				'''wb'''      ) as f:
																f.write(to_bytes(state.opt_state      )      )
										joblib.dump(self.args		,				os.path.join(UpperCAmelCase__		,				'''args.joblib'''      )      )
										joblib.dump(self.data_collator		,				os.path.join(UpperCAmelCase__		,				'''data_collator.joblib'''      )      )
										with open(os.path.join(UpperCAmelCase__		,				'''training_state.json'''      )		,				'''w'''      ) as f:
																json.dump({'''step''': state.step.item()}		,				UpperCAmelCase__      )
										print('''DONE'''      )






def 							_lowerCAmelCase (      __magic_name__	:    Optional[Any]      ,   __magic_name__	:    List[Any]	)							->     Optional[Any]:
						print(f'''RESTORING CHECKPOINT FROM {save_dir}'''      ,   end=''' ... '''	)
						with open(os.path.join(__magic_name__      ,   '''flax_model.msgpack'''	)      ,   '''rb'''	) as f:
												lowercase   :       int								=from_bytes(state.params      ,   f.read()	)

						with open(os.path.join(__magic_name__      ,   '''opt_state.msgpack'''	)      ,   '''rb'''	) as f:
												lowercase   :       Union[str, Any]								=from_bytes(state.opt_state      ,   f.read()	)

						lowercase   :       List[str]								=joblib.load(os.path.join(__magic_name__      ,   '''args.joblib'''	)	)
						lowercase   :       Any								=joblib.load(os.path.join(__magic_name__      ,   '''data_collator.joblib'''	)	)

						with open(os.path.join(__magic_name__      ,   '''training_state.json'''	)      ,   '''r'''	) as f:
												lowercase   :       Any								=json.load(__magic_name__	)
						lowercase   :       Dict								=training_state['''step''']

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






def 							_lowerCAmelCase (      __magic_name__	:    Optional[Any]      ,   __magic_name__	:    Optional[Any]      ,   __magic_name__	:    Optional[Any]      ,   __magic_name__	:    str	)							->     str:
						lowercase   :       str								=num_train_steps - warmup_steps
						lowercase   :       Dict								=optax.linear_schedule(init_value=__magic_name__      ,   end_value=__magic_name__      ,   transition_steps=__magic_name__	)
						lowercase   :       Tuple								=optax.linear_schedule(init_value=__magic_name__      ,   end_value=1E-7      ,   transition_steps=__magic_name__	)
						lowercase   :       List[str]								=optax.join_schedules(schedules=[warmup_fn, decay_fn]      ,   boundaries=[warmup_steps]	)
						return lr







def 							_lowerCAmelCase (      __magic_name__	:    Optional[Any]      ,   __magic_name__	:    List[str]      ,   __magic_name__	:    List[Any]      ,   __magic_name__	:    str      ,   __magic_name__	:    List[str]	)							->     List[Any]:
						def weight_decay_mask(__magic_name__	:    Any	):
												lowercase   :       Tuple								=traverse_util.flatten_dict(__magic_name__	)
												lowercase   :       Dict								={k: (v[-1] != '''bias''' and v[-2:] != ('''LayerNorm''', '''scale''')) for k, v in params.items()}
												return traverse_util.unflatten_dict(__magic_name__	)

						lowercase   :       str								=scheduler_fn(__magic_name__      ,   __magic_name__      ,   __magic_name__      ,   __magic_name__	)

						lowercase   :       str								=optax.adamw(learning_rate=__magic_name__      ,   weight_decay=__magic_name__      ,   mask=__magic_name__	)
						return tx, lr








'''simple docstring'''





import unittest

import numpy as np

from transformers import RobertaConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_flax_available():
							from transformers.models.roberta.modeling_flax_roberta import (
							    FlaxRobertaForCausalLM,
							    FlaxRobertaForMaskedLM,
							    FlaxRobertaForMultipleChoice,
							    FlaxRobertaForQuestionAnswering,
							    FlaxRobertaForSequenceClassification,
							    FlaxRobertaForTokenClassification,
							    FlaxRobertaModel,
							)





class 		__SCREAMING_SNAKE_CASE       ( unittest.TestCase			):







				def __init__(					self    :		Optional[Any]		,				UpperCAmelCase__    :		Tuple		,				UpperCAmelCase__    :		str=13		,				UpperCAmelCase__    :		Optional[Any]=7		,				UpperCAmelCase__    :		Dict=True		,				UpperCAmelCase__    :		int=True		,				UpperCAmelCase__    :		int=True		,				UpperCAmelCase__    :		Dict=True		,				UpperCAmelCase__    :		Dict=99		,				UpperCAmelCase__    :		str=32		,				UpperCAmelCase__    :		Optional[Any]=5		,				UpperCAmelCase__    :		Any=4		,				UpperCAmelCase__    :		List[Any]=37		,				UpperCAmelCase__    :		str="gelu"		,				UpperCAmelCase__    :		Dict=0.1		,				UpperCAmelCase__    :		List[Any]=0.1		,				UpperCAmelCase__    :		Optional[Any]=512		,				UpperCAmelCase__    :		Any=16		,				UpperCAmelCase__    :		str=2		,				UpperCAmelCase__    :		Tuple=0.02		,				UpperCAmelCase__    :		int=4		,				):
										'''simple docstring'''
										lowercase   :       int								=parent
										lowercase   :       List[str]								=batch_size
										lowercase   :       str								=seq_length
										lowercase   :       Optional[Any]								=is_training
										lowercase   :       Union[str, Any]								=use_attention_mask
										lowercase   :       Optional[Any]								=use_token_type_ids
										lowercase   :       Tuple								=use_labels
										lowercase   :       List[str]								=vocab_size
										lowercase   :       List[str]								=hidden_size
										lowercase   :       Tuple								=num_hidden_layers
										lowercase   :       Any								=num_attention_heads
										lowercase   :       List[str]								=intermediate_size
										lowercase   :       Optional[Any]								=hidden_act
										lowercase   :       Dict								=hidden_dropout_prob
										lowercase   :       List[Any]								=attention_probs_dropout_prob
										lowercase   :       Optional[Any]								=max_position_embeddings
										lowercase   :       Tuple								=type_vocab_size
										lowercase   :       Optional[int]								=type_sequence_label_size
										lowercase   :       Optional[Any]								=initializer_range
										lowercase   :       Optional[int]								=num_choices







				def    lowerCamelCase_  (					self    :		str      ):
										'''simple docstring'''
										lowercase   :       Optional[Any]								=ids_tensor([self.batch_size, self.seq_length]		,				self.vocab_size      )

										lowercase   :       Union[str, Any]								=None
										if self.use_attention_mask:
																lowercase   :       List[str]								=random_attention_mask([self.batch_size, self.seq_length]      )

										lowercase   :       Tuple								=None
										if self.use_token_type_ids:
																lowercase   :       str								=ids_tensor([self.batch_size, self.seq_length]		,				self.type_vocab_size      )

										lowercase   :       int								=RobertaConfig(
										    vocab_size=self.vocab_size		,				hidden_size=self.hidden_size		,				num_hidden_layers=self.num_hidden_layers		,				num_attention_heads=self.num_attention_heads		,				intermediate_size=self.intermediate_size		,				hidden_act=self.hidden_act		,				hidden_dropout_prob=self.hidden_dropout_prob		,				attention_probs_dropout_prob=self.attention_probs_dropout_prob		,				max_position_embeddings=self.max_position_embeddings		,				type_vocab_size=self.type_vocab_size		,				is_decoder=UpperCAmelCase__		,				initializer_range=self.initializer_range		,				)

										return config, input_ids, token_type_ids, attention_mask







				def    lowerCamelCase_  (					self    :		Dict      ):
										'''simple docstring'''
										lowercase   :       List[Any]								=self.prepare_config_and_inputs()
										lowercase		,  lowercase		,  lowercase		,  lowercase   :       str								=config_and_inputs
										lowercase   :       Optional[Any]								={'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': attention_mask}
										return config, inputs_dict







				def    lowerCamelCase_  (					self    :		Optional[Any]      ):
										'''simple docstring'''
										lowercase   :       List[str]								=self.prepare_config_and_inputs()
										lowercase		,  lowercase		,  lowercase		,  lowercase   :       Any								=config_and_inputs

										lowercase   :       List[str]								=True
										lowercase   :       Optional[Any]								=floats_tensor([self.batch_size, self.seq_length, self.hidden_size]      )
										lowercase   :       str								=ids_tensor([self.batch_size, self.seq_length]		,				vocab_size=2      )

										return (
										    config,
										    input_ids,
										    token_type_ids,
										    encoder_hidden_states,
										    encoder_attention_mask,
										)





@require_flax
class 		__SCREAMING_SNAKE_CASE       ( lowercase__	,       unittest.TestCase			):
				lowerCamelCase_							=						True

				lowerCamelCase_							=						(
				    (
				        FlaxRobertaModel,
				        FlaxRobertaForCausalLM,
				        FlaxRobertaForMaskedLM,
				        FlaxRobertaForSequenceClassification,
				        FlaxRobertaForTokenClassification,
				        FlaxRobertaForMultipleChoice,
				        FlaxRobertaForQuestionAnswering,
				    )
				    if is_flax_available()
				    else ()
				)







				def    lowerCamelCase_  (					self    :		List[Any]      ):
										'''simple docstring'''
										lowercase   :       str								=FlaxRobertaModelTester(self      )







				@slow
				def    lowerCamelCase_  (					self    :		Optional[int]      ):
										'''simple docstring'''
										for model_class_name in self.all_model_classes:
																lowercase   :       Optional[int]								=model_class_name.from_pretrained('''roberta-base'''		,				from_pt=UpperCAmelCase__      )
																lowercase   :       List[Any]								=model(np.ones((1, 1)      )      )
																self.assertIsNotNone(UpperCAmelCase__      )





