




"""simple docstring"""





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      a__				:


			def __init__(      self      :    Any      ,  UpperCamelCase_      :    Union[str, Any]      ,  UpperCamelCase_      :    List[Any]=13      ,  UpperCamelCase_      :    Tuple=7      ,  UpperCamelCase_      :    List[Any]=True      ,  UpperCamelCase_      :    Any=True      ,  UpperCamelCase_      :    str=True      ,  UpperCamelCase_      :    str=True      ,  UpperCamelCase_      :    Union[str, Any]=True      ,  UpperCamelCase_      :    List[Any]=False      ,  UpperCamelCase_      :    List[Any]=False      ,  UpperCamelCase_      :    List[str]=False      ,  UpperCamelCase_      :    Optional[Any]=2      ,  UpperCamelCase_      :    Any=99      ,  UpperCamelCase_      :    Optional[Any]=0      ,  UpperCamelCase_      :    Any=32      ,  UpperCamelCase_      :    List[Any]=5      ,  UpperCamelCase_      :    List[str]=4      ,  UpperCamelCase_      :    Tuple=0.1      ,  UpperCamelCase_      :    int=0.1      ,  UpperCamelCase_      :    str=512      ,  UpperCamelCase_      :    List[Any]=2      ,  UpperCamelCase_      :    List[Any]=0.02      ,  UpperCamelCase_      :    List[Any]=2      ,  UpperCamelCase_      :    Any=4      ,  UpperCamelCase_      :    str="last"      ,  UpperCamelCase_      :    Optional[int]=True      ,  UpperCamelCase_      :    Any=None      ,  UpperCamelCase_      :    Tuple=0      ,  ):


						"""simple docstring"""




						__UpperCAmelCase	: int		=					parent
						__UpperCAmelCase	: int		=					batch_size
						__UpperCAmelCase	: Optional[Any]		=					seq_length
						__UpperCAmelCase	: str		=					is_training
						__UpperCAmelCase	: str		=					use_input_lengths
						__UpperCAmelCase	: str		=					use_token_type_ids
						__UpperCAmelCase	: List[str]		=					use_labels
						__UpperCAmelCase	: Any		=					gelu_activation
						__UpperCAmelCase	: Optional[Any]		=					sinusoidal_embeddings
						__UpperCAmelCase	: Any		=					causal
						__UpperCAmelCase	: List[str]		=					asm
						__UpperCAmelCase	: Union[str, Any]		=					n_langs
						__UpperCAmelCase	: List[Any]		=					vocab_size
						__UpperCAmelCase	: Dict		=					n_special
						__UpperCAmelCase	: str		=					hidden_size
						__UpperCAmelCase	: List[str]		=					num_hidden_layers
						__UpperCAmelCase	: Optional[Any]		=					num_attention_heads
						__UpperCAmelCase	: int		=					hidden_dropout_prob
						__UpperCAmelCase	: str		=					attention_probs_dropout_prob
						__UpperCAmelCase	: Any		=					max_position_embeddings
						__UpperCAmelCase	: Union[str, Any]		=					type_sequence_label_size
						__UpperCAmelCase	: Any		=					initializer_range
						__UpperCAmelCase	: List[str]		=					num_labels
						__UpperCAmelCase	: str		=					num_choices
						__UpperCAmelCase	: Union[str, Any]		=					summary_type
						__UpperCAmelCase	: Dict		=					use_proj
						__UpperCAmelCase	: Dict		=					scope
						__UpperCAmelCase	: Dict		=					bos_token_id


			def       a_     (      self      :    int):


						"""simple docstring"""




						__UpperCAmelCase	: Optional[Any]		=					ids_tensor([self.batch_size, self.seq_length]      ,  self.vocab_size)
						__UpperCAmelCase	: Optional[int]		=					random_attention_mask([self.batch_size, self.seq_length])

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

						__UpperCAmelCase	: Optional[int]		=					None
						if self.use_token_type_ids:
									__UpperCAmelCase	: Optional[Any]		=					ids_tensor([self.batch_size, self.seq_length]      ,  self.n_langs)

						__UpperCAmelCase	: Dict		=					None
						__UpperCAmelCase	: Union[str, Any]		=					None
						__UpperCAmelCase	: Optional[int]		=					None
						if self.use_labels:
									__UpperCAmelCase	: Optional[int]		=					ids_tensor([self.batch_size]      ,  self.type_sequence_label_size)
									__UpperCAmelCase	: List[str]		=					ids_tensor([self.batch_size, self.seq_length]      ,  self.num_labels)
									__UpperCAmelCase	: int		=					ids_tensor([self.batch_size]      ,  2).float()
									__UpperCAmelCase	: List[Any]		=					ids_tensor([self.batch_size]      ,  self.num_choices)

						__UpperCAmelCase	: List[str]		=					self.get_config()

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


			def       a_     (      self      :    int):


						"""simple docstring"""




						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       a_     (      self      :    List[str]      ,  UpperCamelCase_      :    Optional[int]      ,  UpperCamelCase_      :    int      ,  UpperCamelCase_      :    str      ,  UpperCamelCase_      :    Union[str, Any]      ,  UpperCamelCase_      :    List[str]      ,  UpperCamelCase_      :    str      ,  UpperCamelCase_      :    Tuple      ,  UpperCamelCase_      :    Dict      ,  UpperCamelCase_      :    int      ,  ):


						"""simple docstring"""




						__UpperCAmelCase	: List[Any]		=					XLMModel(config=__lowercase)
						model.to(__lowercase)
						model.eval()
						__UpperCAmelCase	: Optional[int]		=					model(__lowercase      ,  lengths=__lowercase      ,  langs=__lowercase)
						__UpperCAmelCase	: Union[str, Any]		=					model(__lowercase      ,  langs=__lowercase)
						__UpperCAmelCase	: Union[str, Any]		=					model(__lowercase)
						self.parent.assertEqual(result.last_hidden_state.shape      ,  (self.batch_size, self.seq_length, self.hidden_size))


			def       a_     (      self      :    Tuple      ,  UpperCamelCase_      :    Dict      ,  UpperCamelCase_      :    Optional[int]      ,  UpperCamelCase_      :    Tuple      ,  UpperCamelCase_      :    List[Any]      ,  UpperCamelCase_      :    List[Any]      ,  UpperCamelCase_      :    Any      ,  UpperCamelCase_      :    Dict      ,  UpperCamelCase_      :    Optional[Any]      ,  UpperCamelCase_      :    Dict      ,  ):


						"""simple docstring"""




						__UpperCAmelCase	: Optional[Any]		=					XLMWithLMHeadModel(__lowercase)
						model.to(__lowercase)
						model.eval()

						__UpperCAmelCase	: List[str]		=					model(__lowercase      ,  token_type_ids=__lowercase      ,  labels=__lowercase)
						self.parent.assertEqual(result.loss.shape      ,  ())
						self.parent.assertEqual(result.logits.shape      ,  (self.batch_size, self.seq_length, self.vocab_size))


			def       a_     (      self      :    Optional[Any]      ,  UpperCamelCase_      :    Tuple      ,  UpperCamelCase_      :    List[str]      ,  UpperCamelCase_      :    str      ,  UpperCamelCase_      :    List[str]      ,  UpperCamelCase_      :    Tuple      ,  UpperCamelCase_      :    List[str]      ,  UpperCamelCase_      :    Any      ,  UpperCamelCase_      :    Union[str, Any]      ,  UpperCamelCase_      :    int      ,  ):


						"""simple docstring"""




						__UpperCAmelCase	: List[str]		=					XLMForQuestionAnsweringSimple(__lowercase)
						model.to(__lowercase)
						model.eval()

						__UpperCAmelCase	: List[Any]		=					model(__lowercase)

						__UpperCAmelCase	: str		=					model(__lowercase      ,  start_positions=__lowercase      ,  end_positions=__lowercase)
						__UpperCAmelCase	: Union[str, Any]		=					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       a_     (      self      :    Optional[int]      ,  UpperCamelCase_      :    str      ,  UpperCamelCase_      :    Any      ,  UpperCamelCase_      :    Union[str, Any]      ,  UpperCamelCase_      :    Tuple      ,  UpperCamelCase_      :    str      ,  UpperCamelCase_      :    Any      ,  UpperCamelCase_      :    Dict      ,  UpperCamelCase_      :    Union[str, Any]      ,  UpperCamelCase_      :    Any      ,  ):


						"""simple docstring"""




						__UpperCAmelCase	: int		=					XLMForQuestionAnswering(__lowercase)
						model.to(__lowercase)
						model.eval()

						__UpperCAmelCase	: List[str]		=					model(__lowercase)

						__UpperCAmelCase	: Optional[int]		=					model(
						    __lowercase      ,  start_positions=__lowercase      ,  end_positions=__lowercase      ,  cls_index=__lowercase      ,  is_impossible=__lowercase      ,  p_mask=__lowercase      ,  )

						__UpperCAmelCase	: List[str]		=					model(
						    __lowercase      ,  start_positions=__lowercase      ,  end_positions=__lowercase      ,  cls_index=__lowercase      ,  is_impossible=__lowercase      ,  )

						((__UpperCAmelCase       )	,            )	: Any		=					result_with_labels.to_tuple()

						__UpperCAmelCase	: Optional[int]		=					model(__lowercase      ,  start_positions=__lowercase      ,  end_positions=__lowercase)

						((__UpperCAmelCase       )	,            )	: str		=					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       a_     (      self      :    Optional[int]      ,  UpperCamelCase_      :    int      ,  UpperCamelCase_      :    Optional[Any]      ,  UpperCamelCase_      :    Tuple      ,  UpperCamelCase_      :    List[Any]      ,  UpperCamelCase_      :    Tuple      ,  UpperCamelCase_      :    Dict      ,  UpperCamelCase_      :    Dict      ,  UpperCamelCase_      :    Optional[int]      ,  UpperCamelCase_      :    Any      ,  ):


						"""simple docstring"""




						__UpperCAmelCase	: Any		=					XLMForSequenceClassification(__lowercase)
						model.to(__lowercase)
						model.eval()

						__UpperCAmelCase	: int		=					model(__lowercase)
						__UpperCAmelCase	: Dict		=					model(__lowercase      ,  labels=__lowercase)
						self.parent.assertEqual(result.loss.shape      ,  ())
						self.parent.assertEqual(result.logits.shape      ,  (self.batch_size, self.type_sequence_label_size))


			def       a_     (      self      :    Optional[int]      ,  UpperCamelCase_      :    Tuple      ,  UpperCamelCase_      :    Dict      ,  UpperCamelCase_      :    Tuple      ,  UpperCamelCase_      :    Optional[Any]      ,  UpperCamelCase_      :    Union[str, Any]      ,  UpperCamelCase_      :    int      ,  UpperCamelCase_      :    Optional[Any]      ,  UpperCamelCase_      :    int      ,  UpperCamelCase_      :    Optional[int]      ,  ):


						"""simple docstring"""




						__UpperCAmelCase	: int		=					self.num_labels
						__UpperCAmelCase	: Tuple		=					XLMForTokenClassification(__lowercase)
						model.to(__lowercase)
						model.eval()

						__UpperCAmelCase	: Union[str, Any]		=					model(__lowercase      ,  attention_mask=__lowercase      ,  labels=__lowercase)
						self.parent.assertEqual(result.logits.shape      ,  (self.batch_size, self.seq_length, self.num_labels))


			def       a_     (      self      :    List[Any]      ,  UpperCamelCase_      :    Optional[Any]      ,  UpperCamelCase_      :    List[Any]      ,  UpperCamelCase_      :    int      ,  UpperCamelCase_      :    List[Any]      ,  UpperCamelCase_      :    List[str]      ,  UpperCamelCase_      :    Dict      ,  UpperCamelCase_      :    Union[str, Any]      ,  UpperCamelCase_      :    int      ,  UpperCamelCase_      :    Dict      ,  ):


						"""simple docstring"""




						__UpperCAmelCase	: List[Any]		=					self.num_choices
						__UpperCAmelCase	: Tuple		=					XLMForMultipleChoice(config=__lowercase)
						model.to(__lowercase)
						model.eval()
						__UpperCAmelCase	: List[str]		=					input_ids.unsqueeze(1).expand(-1      ,  self.num_choices      ,  -1).contiguous()
						__UpperCAmelCase	: Tuple		=					token_type_ids.unsqueeze(1).expand(-1      ,  self.num_choices      ,  -1).contiguous()
						__UpperCAmelCase	: Optional[Any]		=					input_mask.unsqueeze(1).expand(-1      ,  self.num_choices      ,  -1).contiguous()
						__UpperCAmelCase	: List[Any]		=					model(
						    __lowercase      ,  attention_mask=__lowercase      ,  token_type_ids=__lowercase      ,  labels=__lowercase      ,  )
						self.parent.assertEqual(result.logits.shape      ,  (self.batch_size, self.num_choices))


			def       a_     (      self      :    Optional[int]):


						"""simple docstring"""




						__UpperCAmelCase	: List[Any]		=					self.prepare_config_and_inputs()
						(
						    (
						    __UpperCAmelCase       
						)	,     (
						    __UpperCAmelCase       
						)	,     (
						    __UpperCAmelCase       
						)	,     (
						    __UpperCAmelCase       
						)	,     (
						    __UpperCAmelCase       
						)	,     (
						    __UpperCAmelCase       
						)	,     (
						    __UpperCAmelCase       
						)	,     (
						    __UpperCAmelCase       
						)	,     (
						    __UpperCAmelCase       
						)	,            
						)	: List[Any]		=					config_and_inputs
						__UpperCAmelCase	: Union[str, Any]		=					{"input_ids": input_ids, "token_type_ids": token_type_ids, "lengths": input_lengths}
						return config, inputs_dict






@require_torch
class      a__				(	lowercase__							,    lowercase__							,    lowercase__							,    unittest.TestCase   ):
			lowercase_    =	(
			    (
			        XLMModel,
			        XLMWithLMHeadModel,
			        XLMForQuestionAnswering,
			        XLMForSequenceClassification,
			        XLMForQuestionAnsweringSimple,
			        XLMForTokenClassification,
			        XLMForMultipleChoice,
			    )
			    if is_torch_available()
			    else ()
			)
			lowercase_    =	(
			    (XLMWithLMHeadModel,) if is_torch_available() else ()
			)  # TODO (PVP): Check other models whether language generation is also applicable
			lowercase_    =	(
			    {
			        """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       a_     (      self      :    Tuple      ,  UpperCamelCase_      :    str      ,  UpperCamelCase_      :    Optional[Any]      ,  UpperCamelCase_      :    Any      ,  UpperCamelCase_      :    Optional[Any]      ,  UpperCamelCase_      :    Tuple):


						"""simple docstring"""




						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      :    str      ,  UpperCamelCase_      :    str      ,  UpperCamelCase_      :    Tuple      ,  UpperCamelCase_      :    Optional[int]=False):


						"""simple docstring"""




						__UpperCAmelCase	: str		=					super()._prepare_for_class(__lowercase      ,  __lowercase      ,  return_labels=__lowercase)

						if return_labels:
									if model_class.__name__ == "XLMForQuestionAnswering":
												__UpperCAmelCase	: Optional[Any]		=					torch.zeros(
												    self.model_tester.batch_size      ,  dtype=torch.long      ,  device=__lowercase)
												__UpperCAmelCase	: Optional[int]		=					torch.zeros(
												    self.model_tester.batch_size      ,  dtype=torch.long      ,  device=__lowercase)

						return inputs_dict


			def       a_     (      self      :    Any):


						"""simple docstring"""




						__UpperCAmelCase	: Tuple		=					XLMModelTester(self)
						__UpperCAmelCase	: str		=					ConfigTester(self      ,  config_class=__lowercase      ,  emb_dim=37)


			def       a_     (      self      :    Tuple):


						"""simple docstring"""




						self.config_tester.run_common_tests()


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


						"""simple docstring"""




						__UpperCAmelCase	: List[Any]		=					self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_xlm_model(*__lowercase)


			def       a_     (      self      :    List[str]):


						"""simple docstring"""




						__UpperCAmelCase	: Union[str, Any]		=					self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_xlm_lm_head(*__lowercase)


			def       a_     (      self      :    Tuple):


						"""simple docstring"""




						__UpperCAmelCase	: Dict		=					self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_xlm_simple_qa(*__lowercase)


			def       a_     (      self      :    Any):


						"""simple docstring"""




						__UpperCAmelCase	: Dict		=					self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_xlm_qa(*__lowercase)


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


						"""simple docstring"""




						__UpperCAmelCase	: Any		=					self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_xlm_sequence_classif(*__lowercase)


			def       a_     (      self      :    Optional[int]):


						"""simple docstring"""




						__UpperCAmelCase	: List[Any]		=					self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_xlm_token_classif(*__lowercase)


			def       a_     (      self      :    str):


						"""simple docstring"""




						__UpperCAmelCase	: int		=					self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_xlm_for_multiple_choice(*__lowercase)


			def       a_     (      self      :    Optional[int]      ,  UpperCamelCase_      :    Optional[int]      ,  UpperCamelCase_      :    List[str]      ,  UpperCamelCase_      :    int      ,  UpperCamelCase_      :    str      ,  UpperCamelCase_      :    int      ,  UpperCamelCase_      :    Any=False      ,  UpperCamelCase_      :    Dict=1):


						"""simple docstring"""




						self.assertIsInstance(__lowercase      ,  __lowercase)
						self.assertListEqual(
						    [isinstance(__lowercase      ,  __lowercase) for iter_attentions in attentions]      ,  [True] * len(__lowercase))
						self.assertEqual(len(__lowercase)      ,  (max_length - min_length) * num_beam_groups)

						for idx, iter_attentions in enumerate(__lowercase):
									# adds PAD dummy token
									__UpperCAmelCase	: Union[str, Any]		=					min_length + idx + 1
									__UpperCAmelCase	: List[Any]		=					min_length + idx + 1

									__UpperCAmelCase	: Any		=					(
									    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(__lowercase))


			def       a_     (      self      :    Optional[Any]      ,  UpperCamelCase_      :    str      ,  UpperCamelCase_      :    List[str]      ,  UpperCamelCase_      :    Optional[int]      ,  UpperCamelCase_      :    List[str]      ,  UpperCamelCase_      :    Tuple      ,  UpperCamelCase_      :    Optional[int]=False      ,  UpperCamelCase_      :    int=1):


						"""simple docstring"""




						self.assertIsInstance(__lowercase      ,  __lowercase)
						self.assertListEqual(
						    [isinstance(__lowercase      ,  __lowercase) for iter_hidden_states in hidden_states]      ,  [True] * len(__lowercase)      ,  )
						self.assertEqual(len(__lowercase)      ,  (max_length - min_length) * num_beam_groups)

						for idx, iter_hidden_states in enumerate(__lowercase):
									# adds PAD dummy token
									__UpperCAmelCase	: Tuple		=					min_length + idx + 1
									__UpperCAmelCase	: List[str]		=					(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(__lowercase)      ,  )
						pass


			@slow
			def       a_     (      self      :    Dict):


						"""simple docstring"""




						for model_name in XLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									__UpperCAmelCase	: Union[str, Any]		=					XLMModel.from_pretrained(__lowercase)
									self.assertIsNotNone(__lowercase)






@require_torch
class      a__				(	unittest.TestCase   ):


			@slow
			def       a_     (      self      :    List[str]):


						"""simple docstring"""




						__UpperCAmelCase	: Union[str, Any]		=					XLMWithLMHeadModel.from_pretrained("xlm-mlm-en-2048")
						model.to(__lowercase)
						__UpperCAmelCase	: str		=					torch.tensor([[14, 447]]      ,  dtype=torch.long      ,  device=__lowercase)  # the president
						__UpperCAmelCase	: int		=					[
						    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
						__UpperCAmelCase	: Optional[int]		=					model.generate(__lowercase      ,  do_sample=__lowercase)
						self.assertListEqual(output_ids[0].cpu().numpy().tolist()      ,  __lowercase)








"""simple docstring"""





from typing import Optional

import pyspark

from .. import Features, NamedSplit
from ..download import DownloadMode
from ..packaged_modules.spark.spark import Spark
from .abc import AbstractDatasetReader



class      a__				(	__magic_name__   ):


			def __init__(      self      :    List[Any]      ,  UpperCamelCase_      :    pyspark.sql.DataFrame      ,  UpperCamelCase_      :    Optional[NamedSplit] = None      ,  UpperCamelCase_      :    Optional[Features] = None      ,  UpperCamelCase_      :    bool = True      ,  UpperCamelCase_      :    str = None      ,  UpperCamelCase_      :    bool = False      ,  UpperCamelCase_      :    str = None      ,  UpperCamelCase_      :    bool = True      ,  UpperCamelCase_      :    str = "arrow"      ,  **UpperCamelCase_      :    Tuple      ,  ):


						"""simple docstring"""




						super().__init__(
						    split=UpperCamelCase_      ,  features=UpperCamelCase_      ,  cache_dir=UpperCamelCase_      ,  keep_in_memory=UpperCamelCase_      ,  streaming=UpperCamelCase_      ,  **UpperCamelCase_      ,  )
						__UpperCAmelCase	: List[Any]		=					load_from_cache_file
						__UpperCAmelCase	: List[Any]		=					file_format
						__UpperCAmelCase	: str		=					Spark(
						    df=UpperCamelCase_      ,  features=UpperCamelCase_      ,  cache_dir=UpperCamelCase_      ,  working_dir=UpperCamelCase_      ,  **UpperCamelCase_      ,  )


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


						"""simple docstring"""




						if self.streaming:
									return self.builder.as_streaming_dataset(split=self.split)
						__UpperCAmelCase	: List[Any]		=					None if self._load_from_cache_file else DownloadMode.FORCE_REDOWNLOAD
						self.builder.download_and_prepare(
						    download_mode=UpperCamelCase_      ,  file_format=self._file_format      ,  )
						return self.builder.as_dataset(split=self.split)


