


import argparse
import glob
import logging
import os
from argparse import Namespace
from importlib import import_module

import numpy as np
import torch
from lightning_base import BaseTransformer, add_generic_args, generic_train
from seqeval.metrics import accuracy_score, fa_score, precision_score, recall_score
from torch.nn import CrossEntropyLoss
from torch.utils.data import DataLoader, TensorDataset
from utils_ner import TokenClassificationTask


_lowercase							:   Dict						=						logging.getLogger(__name__)






class    _UpperCamelCase     (       __snake_case					):

				"""simple docstring"""







				lowerCAmelCase					=			'token-classification'

				def __init__(    self					, a__ )						->					List[Any]:
					if type(a__ ) == dict:
						A             =					Namespace(**a__ )
					A             =					import_module("""tasks""" )
					try:
						A             =					getattr(a__					, hparams.task_type )
						A             =					token_classification_task_clazz()
					except AttributeError:
						raise ValueError(
						    f'Task {hparams.task_type} needs to be defined as a TokenClassificationTask subclass in {module}. '
						    f'Available tasks classes are: {TokenClassificationTask.__subclasses__()}' )
					A             =					self.token_classification_task.get_labels(hparams.labels )
					A             =					CrossEntropyLoss().ignore_index
					super().__init__(a__					, len(self.labels )					, self.mode )

				def   _UpperCAmelCase       (    self					, **a__ )						->					List[Any]:
					return self.model(**a__ )

				def   _UpperCAmelCase       (    self					, a__					, a__ )						->					Dict:
					A             =					{"""input_ids""": batch[0], """attention_mask""": batch[1], """labels""": batch[3]}
					if self.config.model_type != "distilbert":
						A             =					(
						    batch[2] if self.config.model_type in ["""bert""", """xlnet"""] else None
						)  # XLM and RoBERTa don"t use token_type_ids

					A             =					self(**a__ )
					A             =					outputs[0]
					# tensorboard_logs = {"loss": loss, "rate": self.lr_scheduler.get_last_lr()[-1]}
					return {"loss": loss}

				def   _UpperCAmelCase       (    self )						->					str:
					A             =					self.hparams
					for mode in ["train", "dev", "test"]:
						A             =					self._feature_file(a__ )
						if os.path.exists(a__ ) and not args.overwrite_cache:
							logger.info("""Loading features from cached file %s"""					, a__ )
							A             =					torch.load(a__ )
						else:
							logger.info("""Creating features from dataset file at %s"""					, args.data_dir )
							A             =					self.token_classification_task.read_examples_from_file(args.data_dir					, a__ )
							A             =					self.token_classification_task.convert_examples_to_features(
							    a__					, self.labels					, args.max_seq_length					, self.tokenizer					, cls_token_at_end=bool(self.config.model_type in ["""xlnet"""] )					, cls_token=self.tokenizer.cls_token					, cls_token_segment_id=2 if self.config.model_type in ["""xlnet"""] else 0					, sep_token=self.tokenizer.sep_token					, sep_token_extra=a__					, pad_on_left=bool(self.config.model_type in ["""xlnet"""] )					, pad_token=self.tokenizer.pad_token_id					, pad_token_segment_id=self.tokenizer.pad_token_type_id					, pad_token_label_id=self.pad_token_label_id					, )
							logger.info("""Saving features into cached file %s"""					, a__ )
							torch.save(a__					, a__ )

				def   _UpperCAmelCase       (    self					, a__					, a__					, a__ = False )						->					DataLoader:
					A             =					self._feature_file(a__ )
					logger.info("""Loading features from cached file %s"""					, a__ )
					A             =					torch.load(a__ )
					A             =					torch.tensor([f.input_ids for f in features]					, dtype=torch.long )
					A             =					torch.tensor([f.attention_mask for f in features]					, dtype=torch.long )
					if features[0].token_type_ids is not None:
						A             =					torch.tensor([f.token_type_ids for f in features]					, dtype=torch.long )
					else:
						A             =					torch.tensor([0 for f in features]					, dtype=torch.long )
						# HACK(we will not use this anymore soon)
					A             =					torch.tensor([f.label_ids for f in features]					, dtype=torch.long )
					return DataLoader(
					    TensorDataset(a__					, a__					, a__					, a__ )					, batch_size=a__ )

				def   _UpperCAmelCase       (    self					, a__					, a__ )						->					Any:
					"""Compute validation""" ""
					A             =					{"""input_ids""": batch[0], """attention_mask""": batch[1], """labels""": batch[3]}
					if self.config.model_type != "distilbert":
						A             =					(
						    batch[2] if self.config.model_type in ["""bert""", """xlnet"""] else None
						)  # XLM and RoBERTa don"t use token_type_ids
					A             =					self(**a__ )
					A       ,      A             =					outputs[:2]
					A             =					logits.detach().cpu().numpy()
					A             =					inputs["""labels"""].detach().cpu().numpy()
					return {"val_loss": tmp_eval_loss.detach().cpu(), "pred": preds, "target": out_label_ids}

				def   _UpperCAmelCase       (    self					, a__ )						->					Optional[int]:
					A             =					torch.stack([x["""val_loss"""] for x in outputs] ).mean()
					A             =					np.concatenate([x["""pred"""] for x in outputs]					, axis=0 )
					A             =					np.argmax(a__					, axis=2 )
					A             =					np.concatenate([x["""target"""] for x in outputs]					, axis=0 )

					A             =					dict(enumerate(self.labels ) )
					A             =					[[] for _ in range(out_label_ids.shape[0] )]
					A             =					[[] for _ in range(out_label_ids.shape[0] )]

					for i in range(out_label_ids.shape[0] ):
						for j in range(out_label_ids.shape[1] ):
							if out_label_ids[i, j] != self.pad_token_label_id:
								out_label_list[i].append(label_map[out_label_ids[i][j]] )
								preds_list[i].append(label_map[preds[i][j]] )

					A             =					{
					    """val_loss""": val_loss_mean,
					    """accuracy_score""": accuracy_score(a__					, a__ ),
					    """precision""": precision_score(a__					, a__ ),
					    """recall""": recall_score(a__					, a__ ),
					    """f1""": fa_score(a__					, a__ ),
					}

					A             =					dict(results.items() )
					A             =					results
					return ret, preds_list, out_label_list

				def   _UpperCAmelCase       (    self					, a__ )						->					Any:
					# when stable
					A       ,      A       ,      A             =					self._eval_end(a__ )
					A             =					ret["""log"""]
					return {"val_loss": logs["val_loss"], "log": logs, "progress_bar": logs}

				def   _UpperCAmelCase       (    self					, a__ )						->					List[Any]:
					# updating to test_epoch_end instead of deprecated test_end
					A       ,      A       ,      A             =					self._eval_end(a__ )

					# Converting to the dict required by pl
					# https://github.com/PyTorchLightning/pytorch-lightning/blob/master/\
					# pytorch_lightning/trainer/logging.py#L139
					A             =					ret["""log"""]
					# `val_loss` is the key returned by `self._eval_end()` but actually refers to `test_loss`
					return {"avg_test_loss": logs["val_loss"], "log": logs, "progress_bar": logs}

				@staticmethod
				def   _UpperCAmelCase       (    a__					, a__ )						->					Dict:
					# Add NER specific options
					BaseTransformer.add_model_specific_args(a__					, a__ )
					parser.add_argument(
					    """--task_type"""					, default="""NER"""					, type=a__					, help="""Task type to fine tune in training (e.g. NER, POS, etc)""" )
					parser.add_argument(
					    """--max_seq_length"""					, default=128					, type=a__					, help=(
					        """The maximum total input sequence length after tokenization. Sequences longer """
					        """than this will be truncated, sequences shorter will be padded."""
					    )					, )

					parser.add_argument(
					    """--labels"""					, default=""""""					, type=a__					, help="""Path to a file containing all labels. If not specified, CoNLL-2003 labels are used."""					, )
					parser.add_argument(
					    """--gpus"""					, default=0					, type=a__					, help="""The number of GPUs allocated for this, it is by default 0 meaning none"""					, )

					parser.add_argument(
					    """--overwrite_cache"""					, action="""store_true"""					, help="""Overwrite the cached training and evaluation sets""" )

					return parser


if __name__ == "__main__":
				_lowercase							:   Optional[Any]						=						argparse.ArgumentParser()
				add_generic_args(parser, os.getcwd())
				_lowercase							:   Dict						=						NERTransformer.add_model_specific_args(parser, os.getcwd())
				_lowercase							:   Union[str, Any]						=						parser.parse_args()
				_lowercase							:   int						=						NERTransformer(args)
				_lowercase							:   Union[str, Any]						=						generic_train(model, args)

				if args.do_predict:
								# See https://github.com/huggingface/transformers/issues/3159
								# pl use this default format to create a checkpoint:
								# https://github.com/PyTorchLightning/pytorch-lightning/blob/master\
								# /pytorch_lightning/callbacks/model_checkpoint.py#L322
								_lowercase							:   Union[str, Any]						=						sorted(glob.glob(os.path.join(args.output_dir, "checkpoint-epoch=*.ckpt"), recursive=True))
								_lowercase							:   Union[str, Any]						=						model.load_from_checkpoint(checkpoints[-1])
								trainer.test(model)







import unittest

from transformers import XLMConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
				import torch

				from transformers import (
				    XLMForMultipleChoice,
				    XLMForQuestionAnswering,
				    XLMForQuestionAnsweringSimple,
				    XLMForSequenceClassification,
				    XLMForTokenClassification,
				    XLMModel,
				    XLMWithLMHeadModel,
				)
				from transformers.models.xlm.modeling_xlm import XLM_PRETRAINED_MODEL_ARCHIVE_LIST






class    _UpperCamelCase     :

				"""simple docstring"""







				def __init__(    self					, a__					, a__=13					, a__=7					, a__=True					, a__=True					, a__=True					, a__=True					, a__=True					, a__=False					, a__=False					, a__=False					, a__=2					, a__=99					, a__=0					, a__=32					, a__=5					, a__=4					, a__=0.1					, a__=0.1					, a__=512					, a__=2					, a__=0.02					, a__=2					, a__=4					, a__="last"					, a__=True					, a__=None					, a__=0					, )						->					Any:
					A             =					parent
					A             =					batch_size
					A             =					seq_length
					A             =					is_training
					A             =					use_input_lengths
					A             =					use_token_type_ids
					A             =					use_labels
					A             =					gelu_activation
					A             =					sinusoidal_embeddings
					A             =					causal
					A             =					asm
					A             =					n_langs
					A             =					vocab_size
					A             =					n_special
					A             =					hidden_size
					A             =					num_hidden_layers
					A             =					num_attention_heads
					A             =					hidden_dropout_prob
					A             =					attention_probs_dropout_prob
					A             =					max_position_embeddings
					A             =					type_sequence_label_size
					A             =					initializer_range
					A             =					num_labels
					A             =					num_choices
					A             =					summary_type
					A             =					use_proj
					A             =					scope
					A             =					bos_token_id

				def   _UpperCAmelCase       (    self )						->					Optional[int]:
					A             =					ids_tensor([self.batch_size, self.seq_length]					, self.vocab_size )
					A             =					random_attention_mask([self.batch_size, self.seq_length] )

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

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

					A             =					None
					A             =					None
					A             =					None
					if self.use_labels:
						A             =					ids_tensor([self.batch_size]					, self.type_sequence_label_size )
						A             =					ids_tensor([self.batch_size, self.seq_length]					, self.num_labels )
						A             =					ids_tensor([self.batch_size]					, 2 ).float()
						A             =					ids_tensor([self.batch_size]					, self.num_choices )

					A             =					self.get_config()

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

				def   _UpperCAmelCase       (    self )						->					int:
					return 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					, num_labels=self.num_labels					, bos_token_id=self.bos_token_id					, )

				def   _UpperCAmelCase       (    self					, a__					, a__					, a__					, a__					, a__					, a__					, a__					, a__					, a__					, )						->					str:
					A             =					XLMModel(config=a__ )
					model.to(a__ )
					model.eval()
					A             =					model(a__					, lengths=a__					, langs=a__ )
					A             =					model(a__					, langs=a__ )
					A             =					model(a__ )
					self.parent.assertEqual(result.last_hidden_state.shape					, (self.batch_size, self.seq_length, self.hidden_size) )

				def   _UpperCAmelCase       (    self					, a__					, a__					, a__					, a__					, a__					, a__					, a__					, a__					, a__					, )						->					str:
					A             =					XLMWithLMHeadModel(a__ )
					model.to(a__ )
					model.eval()

					A             =					model(a__					, token_type_ids=a__					, labels=a__ )
					self.parent.assertEqual(result.loss.shape					, () )
					self.parent.assertEqual(result.logits.shape					, (self.batch_size, self.seq_length, self.vocab_size) )

				def   _UpperCAmelCase       (    self					, a__					, a__					, a__					, a__					, a__					, a__					, a__					, a__					, a__					, )						->					Any:
					A             =					XLMForQuestionAnsweringSimple(a__ )
					model.to(a__ )
					model.eval()

					A             =					model(a__ )

					A             =					model(a__					, start_positions=a__					, end_positions=a__ )
					A             =					outputs
					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   _UpperCAmelCase       (    self					, a__					, a__					, a__					, a__					, a__					, a__					, a__					, a__					, a__					, )						->					Optional[Any]:
					A             =					XLMForQuestionAnswering(a__ )
					model.to(a__ )
					model.eval()

					A             =					model(a__ )

					A             =					model(
					    a__					, start_positions=a__					, end_positions=a__					, cls_index=a__					, is_impossible=a__					, p_mask=a__					, )

					A             =					model(
					    a__					, start_positions=a__					, end_positions=a__					, cls_index=a__					, is_impossible=a__					, )

					((A)       ,      )             =					result_with_labels.to_tuple()

					A             =					model(a__					, start_positions=a__					, end_positions=a__ )

					((A)       ,      )             =					result_with_labels.to_tuple()

					self.parent.assertEqual(result_with_labels.loss.shape					, () )
					self.parent.assertEqual(result.start_top_log_probs.shape					, (self.batch_size, model.config.start_n_top) )
					self.parent.assertEqual(result.start_top_index.shape					, (self.batch_size, model.config.start_n_top) )
					self.parent.assertEqual(
					    result.end_top_log_probs.shape					, (self.batch_size, model.config.start_n_top * model.config.end_n_top) )
					self.parent.assertEqual(
					    result.end_top_index.shape					, (self.batch_size, model.config.start_n_top * model.config.end_n_top) )
					self.parent.assertEqual(result.cls_logits.shape					, (self.batch_size,) )

				def   _UpperCAmelCase       (    self					, a__					, a__					, a__					, a__					, a__					, a__					, a__					, a__					, a__					, )						->					Optional[Any]:
					A             =					XLMForSequenceClassification(a__ )
					model.to(a__ )
					model.eval()

					A             =					model(a__ )
					A             =					model(a__					, labels=a__ )
					self.parent.assertEqual(result.loss.shape					, () )
					self.parent.assertEqual(result.logits.shape					, (self.batch_size, self.type_sequence_label_size) )

				def   _UpperCAmelCase       (    self					, a__					, a__					, a__					, a__					, a__					, a__					, a__					, a__					, a__					, )						->					List[str]:
					A             =					self.num_labels
					A             =					XLMForTokenClassification(a__ )
					model.to(a__ )
					model.eval()

					A             =					model(a__					, attention_mask=a__					, labels=a__ )
					self.parent.assertEqual(result.logits.shape					, (self.batch_size, self.seq_length, self.num_labels) )

				def   _UpperCAmelCase       (    self					, a__					, a__					, a__					, a__					, a__					, a__					, a__					, a__					, a__					, )						->					Dict:
					A             =					self.num_choices
					A             =					XLMForMultipleChoice(config=a__ )
					model.to(a__ )
					model.eval()
					A             =					input_ids.unsqueeze(1 ).expand(-1					, self.num_choices					, -1 ).contiguous()
					A             =					token_type_ids.unsqueeze(1 ).expand(-1					, self.num_choices					, -1 ).contiguous()
					A             =					input_mask.unsqueeze(1 ).expand(-1					, self.num_choices					, -1 ).contiguous()
					A             =					model(
					    a__					, attention_mask=a__					, token_type_ids=a__					, labels=a__					, )
					self.parent.assertEqual(result.logits.shape					, (self.batch_size, self.num_choices) )

				def   _UpperCAmelCase       (    self )						->					Optional[int]:
					A             =					self.prepare_config_and_inputs()
					(
					    (
					    A
					)       ,      (
					    A
					)       ,      (
					    A
					)       ,      (
					    A
					)       ,      (
					    A
					)       ,      (
					    A
					)       ,      (
					    A
					)       ,      (
					    A
					)       ,      (
					    A
					)       ,      
					)             =					config_and_inputs
					A             =					{"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """lengths""": input_lengths}
					return config, inputs_dict






@require_torch
class    _UpperCamelCase     (       __snake_case     ,		__snake_case     ,		__snake_case     ,		unittest.TestCase					):

				"""simple docstring"""







				lowerCAmelCase					=			(
				    (
				        XLMModel,
				        XLMWithLMHeadModel,
				        XLMForQuestionAnswering,
				        XLMForSequenceClassification,
				        XLMForQuestionAnsweringSimple,
				        XLMForTokenClassification,
				        XLMForMultipleChoice,
				    )
				    if is_torch_available()
				    else ()
				)
				lowerCAmelCase					=			(
				    (XLMWithLMHeadModel,) if is_torch_available() else ()
				)  # TODO (PVP): Check other models whether language generation is also applicable
				lowerCAmelCase					=			(
				    {
				        'feature-extraction': XLMModel,
				        'fill-mask': XLMWithLMHeadModel,
				        'question-answering': XLMForQuestionAnsweringSimple,
				        'text-classification': XLMForSequenceClassification,
				        'text-generation': XLMWithLMHeadModel,
				        'token-classification': XLMForTokenClassification,
				        'zero-shot': XLMForSequenceClassification,
				    }
				    if is_torch_available()
				    else {}
				)

				def   _UpperCAmelCase       (    self					, a__					, a__					, a__					, a__					, a__ )						->					Union[str, Any]:
					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   _UpperCAmelCase       (    self					, a__					, a__					, a__=False )						->					Optional[Any]:
					A             =					super()._prepare_for_class(a__					, a__					, return_labels=a__ )

					if return_labels:
						if model_class.__name__ == "XLMForQuestionAnswering":
							A             =					torch.zeros(
							    self.model_tester.batch_size					, dtype=torch.long					, device=a__ )
							A             =					torch.zeros(
							    self.model_tester.batch_size					, dtype=torch.long					, device=a__ )

					return inputs_dict

				def   _UpperCAmelCase       (    self )						->					Union[str, Any]:
					A             =					XLMModelTester(self )
					A             =					ConfigTester(self					, config_class=a__					, emb_dim=37 )

				def   _UpperCAmelCase       (    self )						->					str:
					self.config_tester.run_common_tests()

				def   _UpperCAmelCase       (    self )						->					Any:
					A             =					self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_xlm_model(*a__ )

				def   _UpperCAmelCase       (    self )						->					Optional[int]:
					A             =					self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_xlm_lm_head(*a__ )

				def   _UpperCAmelCase       (    self )						->					Dict:
					A             =					self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_xlm_simple_qa(*a__ )

				def   _UpperCAmelCase       (    self )						->					Optional[Any]:
					A             =					self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_xlm_qa(*a__ )

				def   _UpperCAmelCase       (    self )						->					Optional[Any]:
					A             =					self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_xlm_sequence_classif(*a__ )

				def   _UpperCAmelCase       (    self )						->					Optional[Any]:
					A             =					self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_xlm_token_classif(*a__ )

				def   _UpperCAmelCase       (    self )						->					Tuple:
					A             =					self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_xlm_for_multiple_choice(*a__ )

				def   _UpperCAmelCase       (    self					, a__					, a__					, a__					, a__					, a__					, a__=False					, a__=1 )						->					List[str]:
					self.assertIsInstance(a__					, a__ )
					self.assertListEqual(
					    [isinstance(a__					, a__ ) for iter_attentions in attentions]					, [True] * len(a__ ) )
					self.assertEqual(len(a__ )					, (max_length - min_length) * num_beam_groups )

					for idx, iter_attentions in enumerate(a__ ):
						# adds PAD dummy token
						A             =					min_length + idx + 1
						A             =					min_length + idx + 1

						A             =					(
						    batch_size * num_beam_groups,
						    config.num_attention_heads,
						    tgt_len,
						    src_len,
						)
						# check attn size
						self.assertListEqual(
						    [layer_attention.shape for layer_attention in iter_attentions]					, [expected_shape] * len(a__ ) )

				def   _UpperCAmelCase       (    self					, a__					, a__					, a__					, a__					, a__					, a__=False					, a__=1 )						->					List[Any]:
					self.assertIsInstance(a__					, a__ )
					self.assertListEqual(
					    [isinstance(a__					, a__ ) for iter_hidden_states in hidden_states]					, [True] * len(a__ )					, )
					self.assertEqual(len(a__ )					, (max_length - min_length) * num_beam_groups )

					for idx, iter_hidden_states in enumerate(a__ ):
						# adds PAD dummy token
						A             =					min_length + idx + 1
						A             =					(batch_size * num_beam_groups, seq_len, config.hidden_size)
						# check hidden size
						self.assertListEqual(
						    [layer_hidden_states.shape for layer_hidden_states in iter_hidden_states]					, [expected_shape] * len(a__ )					, )
					pass

				@slow
				def   _UpperCAmelCase       (    self )						->					Any:
					for model_name in XLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
						A             =					XLMModel.from_pretrained(a__ )
						self.assertIsNotNone(a__ )






@require_torch
class    _UpperCamelCase     (       unittest.TestCase					):

				"""simple docstring"""







				@slow
				def   _UpperCAmelCase       (    self )						->					List[Any]:
					A             =					XLMWithLMHeadModel.from_pretrained("""xlm-mlm-en-2048""" )
					model.to(a__ )
					A             =					torch.tensor([[14, 447]]					, dtype=torch.long					, device=a__ )  # the president
					A             =					[
					    14,
					    447,
					    14,
					    447,
					    14,
					    447,
					    14,
					    447,
					    14,
					    447,
					    14,
					    447,
					    14,
					    447,
					    14,
					    447,
					    14,
					    447,
					    14,
					    447,
					]  # 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
					A             =					model.generate(a__					, do_sample=a__ )
					self.assertListEqual(output_ids[0].cpu().numpy().tolist()					, a__ )



