
'''simple docstring'''







import logging
import math
from functools import partial
from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Tuple, Union

import torch

from .tensor_utils import tensor_tree_map, tree_map







def        SCREAMING_SNAKE_CASE__			(     __A			) ->			List[Tuple[int, ...]]:
	_snake_case     		=	[]
	if isinstance(__A    ,       __A			):
		for v in tree.values():
			shapes.extend(_fetch_dims(__A			)			)
	elif isinstance(__A    ,       (list, tuple)			):
		for t in tree:
			shapes.extend(_fetch_dims(__A			)			)
	elif isinstance(__A    ,       torch.Tensor			):
		shapes.append(tree.shape			)
	else:
		raise ValueError('Not supported'			)

	return shapes







@torch.jit.ignore
def        SCREAMING_SNAKE_CASE__			(     __A    ,       __A			) ->			Tuple[int, ...]:
	_snake_case     		=	[]
	for d in reversed(__A			):
		idx.append(flat_idx % d			)
		_snake_case     		=	flat_idx // d

	return tuple(reversed(__A			)			)







@torch.jit.ignore
def        SCREAMING_SNAKE_CASE__			(     __A    ,       __A    ,       __A    ,       __A = None    ,       __A = None    ,       ) ->			List[Tuple[slice, ...]]:

	# start_edges and end_edges both indicate whether, starting from any given
	# dimension, the start/end index is at the top/bottom edge of the
	# corresponding tensor, modeled as a tree
	def reduce_edge_list(__A			) -> None:
		_snake_case     		=	True
		for i in range(len(__A			)			):
			_snake_case     		=	-1 * (i + 1)
			l[reversed_idx] &= tally
			_snake_case     		=	l[reversed_idx]

	if start_edges is None:
		_snake_case     		=	[s == 0 for s in start]
		reduce_edge_list(__A			)
	if end_edges is None:
		_snake_case     		=	[e == (d - 1) for e, d in zip(__A    ,       __A			)]
		reduce_edge_list(__A			)

	# Base cases. Either start/end are empty and we're done, or the final,
	# one-dimensional tensor can be simply sliced
	if len(__A			) == 0:
		return [()]
	elif len(__A			) == 1:
		return [(slice(start[0]    ,       end[0] + 1			),)]

	_snake_case     		=	[]
	_snake_case     		=	[]

	# Dimensions common to start and end can be selected directly
	for s, e in zip(__A    ,       __A			):
		if s == e:
			path_list.append(slice(__A    ,       s + 1			)			)
		else:
			break

	_snake_case     		=	tuple(__A			)
	_snake_case     		=	len(__A			)

	# start == end, and we're done
	if divergence_idx == len(__A			):
		return [path]

	def upper() -> Tuple[Tuple[slice, ...], ...]:
		assert start_edges is not None
		assert end_edges is not None

		_snake_case     		=	start[divergence_idx]
		return tuple(
		    path + (slice(__A    ,       sdi + 1			),) + s
		    for s in _get_minimal_slice_set(
		        start[divergence_idx + 1 :]    ,       [d - 1 for d in dims[divergence_idx + 1 :]]    ,       dims[divergence_idx + 1 :]    ,       start_edges=start_edges[divergence_idx + 1 :]    ,       end_edges=[True for _ in end_edges[divergence_idx + 1 :]]    ,       )			)

	def lower() -> Tuple[Tuple[slice, ...], ...]:
		assert start_edges is not None
		assert end_edges is not None

		_snake_case     		=	end[divergence_idx]
		return tuple(
		    path + (slice(__A    ,       edi + 1			),) + s
		    for s in _get_minimal_slice_set(
		        [0 for _ in start[divergence_idx + 1 :]]    ,       end[divergence_idx + 1 :]    ,       dims[divergence_idx + 1 :]    ,       start_edges=[True for _ in start_edges[divergence_idx + 1 :]]    ,       end_edges=end_edges[divergence_idx + 1 :]    ,       )			)

	# If both start and end are at the edges of the subtree rooted at
	# divergence_idx, we can just select the whole subtree at once
	if start_edges[divergence_idx] and end_edges[divergence_idx]:
		slices.append(path + (slice(start[divergence_idx]    ,       end[divergence_idx] + 1			),)			)
	# If just start is at the edge, we can grab almost all of the subtree,
	# treating only the ragged bottom edge as an edge case
	elif start_edges[divergence_idx]:
		slices.append(path + (slice(start[divergence_idx]    ,       end[divergence_idx]			),)			)
		slices.extend(lower()			)
	# Analogous to the previous case, but the top is ragged this time
	elif end_edges[divergence_idx]:
		slices.extend(upper()			)
		slices.append(path + (slice(start[divergence_idx] + 1    ,       end[divergence_idx] + 1			),)			)
	# If both sides of the range are ragged, we need to handle both sides
	# separately. If there's contiguous meat in between them, we can index it
	# in one big chunk
	else:
		slices.extend(upper()			)
		_snake_case     		=	end[divergence_idx] - start[divergence_idx]
		if middle_ground > 1:
			slices.append(path + (slice(start[divergence_idx] + 1    ,       end[divergence_idx]			),)			)
		slices.extend(lower()			)

	return slices







@torch.jit.ignore
def        SCREAMING_SNAKE_CASE__			(     __A    ,       __A    ,       __A    ,       __A			) ->			torch.Tensor:
	_snake_case     		=	t.shape[:no_batch_dims]
	_snake_case     		=	list(_flat_idx_to_idx(__A    ,       __A			)			)
	# _get_minimal_slice_set is inclusive
	_snake_case     		=	list(_flat_idx_to_idx(flat_end - 1    ,       __A			)			)

	# Get an ordered list of slices to perform
	_snake_case     		=	_get_minimal_slice_set(
	    __A    ,       __A    ,       __A    ,       )

	_snake_case     		=	[t[s] for s in slices]

	return torch.cat([s.view((-1,) + t.shape[no_batch_dims:]			) for s in sliced_tensors]			)







def        SCREAMING_SNAKE_CASE__			(     __A    ,       __A    ,       __A    ,       __A    ,       __A = False    ,       __A = None    ,       __A = False    ,       ) ->			Any:
	if not (len(__A			) > 0):
		raise ValueError('Must provide at least one input'			)

	_snake_case     		=	[shape[:no_batch_dims] for shape in _fetch_dims(__A			)]
	_snake_case     		=	tuple([max(__A			) for s in zip(*__A			)]			)

	def _prep_inputs(__A			) -> torch.Tensor:
		if not low_mem:
			if not sum(t.shape[:no_batch_dims]			) == no_batch_dims:
				_snake_case     		=	t.expand(orig_batch_dims + t.shape[no_batch_dims:]			)
			_snake_case     		=	t.reshape(-1    ,       *t.shape[no_batch_dims:]			)
		else:
			_snake_case     		=	t.expand(orig_batch_dims + t.shape[no_batch_dims:]			)
		return t

	_snake_case     		=	tensor_tree_map(_prep_inputs    ,       __A			)
	_snake_case     		=	None
	if _out is not None:
		_snake_case     		=	tensor_tree_map(lambda __A			: t.view([-1] + list(t.shape[no_batch_dims:]			)			)    ,       _out			)

	_snake_case     		=	1
	for d in orig_batch_dims:
		flat_batch_dim *= d

	_snake_case     		=	flat_batch_dim // chunk_size + (flat_batch_dim % chunk_size != 0)

	def _select_chunk(__A			) -> torch.Tensor:
		return t[i : i + chunk_size] if t.shape[0] != 1 else t

	_snake_case     		=	0
	_snake_case     		=	prepped_outputs
	for _ in range(__A			):
		# Chunk the input
		if not low_mem:
			_snake_case     		=	_select_chunk
		else:
			_snake_case     		=	partial(
			    _chunk_slice    ,       flat_start=__A    ,       flat_end=min(__A    ,       i + chunk_size			)    ,       no_batch_dims=len(__A			)    ,       )

		_snake_case     		=	tensor_tree_map(__A    ,       __A			)

		# Run the layer on the chunk
		_snake_case     		=	layer(**__A			)

		# Allocate space for the output
		if out is None:
			_snake_case     		=	tensor_tree_map(lambda __A			: t.new_zeros((flat_batch_dim,) + t.shape[1:]			)    ,       __A			)

		# Put the chunk in its pre-allocated space
		if isinstance(__A    ,       __A			):

			def assign(__A    ,       __A			) -> None:
				for k, v in da.items():
					if isinstance(__A    ,       __A			):
						assign(__A    ,       da[k]			)
					else:
						if _add_into_out:
							v[i : i + chunk_size] += da[k]
						else:
							_snake_case     		=	da[k]

			assign(__A    ,       __A			)
		elif isinstance(__A    ,       __A			):
			for xa, xa in zip(__A    ,       __A			):
				if _add_into_out:
					xa[i : i + chunk_size] += xa
				else:
					_snake_case     		=	xa
		elif isinstance(__A    ,       torch.Tensor			):
			if _add_into_out:
				out[i : i + chunk_size] += output_chunk
			else:
				_snake_case     		=	output_chunk
		else:
			raise ValueError('Not supported'			)

		i += chunk_size

	_snake_case     		=	tensor_tree_map(lambda __A			: t.view(orig_batch_dims + t.shape[1:]			)    ,       __A			)

	return out





class 		__UpperCAmelCase   :


							def __init__( self					,     lowerCAmelCase_ = 5_12					,     ):





								"""simple docstring"""






								_snake_case     		=	max_chunk_size
								_snake_case     		=	None
								_snake_case     		=	None


							def 				lowerCamelCase       ( self					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_						):





								"""simple docstring"""






								logging.info('Tuning chunk size...'						)

								if min_chunk_size >= self.max_chunk_size:
									return min_chunk_size

								_snake_case     		=	[2**l for l in range(int(math.log(self.max_chunk_size					,     2						)						) + 1						)]
								_snake_case     		=	[c for c in candidates if c > min_chunk_size]
								_snake_case     		=	[min_chunk_size] + candidates
								candidates[-1] += 4

								def test_chunk_size(lowerCAmelCase_						) -> bool:
									try:
										with torch.no_grad():
											fn(*lowerCAmelCase_					,     chunk_size=lowerCAmelCase_						)
										return True
									except RuntimeError:
										return False

								_snake_case     		=	0
								_snake_case     		=	len(lowerCAmelCase_						) - 1
								while i > min_viable_chunk_size_index:
									_snake_case     		=	test_chunk_size(candidates[i]						)
									if not viable:
										_snake_case     		=	(min_viable_chunk_size_index + i) // 2
									else:
										_snake_case     		=	i
										_snake_case     		=	(i + len(lowerCAmelCase_						) - 1) // 2

								return candidates[min_viable_chunk_size_index]


							def 				lowerCamelCase       ( self					,     lowerCAmelCase_					,     lowerCAmelCase_						):





								"""simple docstring"""






								_snake_case     		=	True
								for aa, aa in zip(lowerCAmelCase_					,     lowerCAmelCase_						):
									assert type(lowerCAmelCase_						) == type(lowerCAmelCase_						)
									if isinstance(lowerCAmelCase_					,     (list, tuple)						):
										consistent &= self._compare_arg_caches(lowerCAmelCase_					,     lowerCAmelCase_						)
									elif isinstance(lowerCAmelCase_					,     lowerCAmelCase_						):
										_snake_case     		=	[v for _, v in sorted(aa.items()					,     key=lambda lowerCAmelCase_						: x[0]						)]
										_snake_case     		=	[v for _, v in sorted(aa.items()					,     key=lambda lowerCAmelCase_						: x[0]						)]
										consistent &= self._compare_arg_caches(lowerCAmelCase_					,     lowerCAmelCase_						)
									else:
										consistent &= aa == aa

								return consistent




							def 				lowerCamelCase       ( self					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     ):





								"""simple docstring"""






								_snake_case     		=	True
								_snake_case     		=	tree_map(lambda lowerCAmelCase_						: a.shape if isinstance(lowerCAmelCase_					,     torch.Tensor						) else a					,     lowerCAmelCase_					,     lowerCAmelCase_						)
								if self.cached_arg_data is not None:
									# If args have changed shape/value, we need to re-tune
									assert len(self.cached_arg_data						) == len(lowerCAmelCase_						)
									_snake_case     		=	self._compare_arg_caches(self.cached_arg_data					,     lowerCAmelCase_						)
								else:
									# Otherwise, we can reuse the precomputed value
									_snake_case     		=	False

								if not consistent:
									_snake_case     		=	self._determine_favorable_chunk_size(
									    lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     )
									_snake_case     		=	arg_data

								assert self.cached_chunk_size is not None

								return self.cached_chunk_size



'''simple docstring'''







import unittest

from transformers import 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
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
					import torch

					from transformers import (
					    OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST,
					    OpenAIGPTConfig,
					    OpenAIGPTDoubleHeadsModel,
					    OpenAIGPTForSequenceClassification,
					    OpenAIGPTLMHeadModel,
					    OpenAIGPTModel,
					)


class 		__UpperCAmelCase   :


							def __init__( self					,     lowerCAmelCase_					,     lowerCAmelCase_=13					,     lowerCAmelCase_=7					,     lowerCAmelCase_=True					,     lowerCAmelCase_=True					,     lowerCAmelCase_=True					,     lowerCAmelCase_=99					,     lowerCAmelCase_=32					,     lowerCAmelCase_=5					,     lowerCAmelCase_=4					,     lowerCAmelCase_=37					,     lowerCAmelCase_="gelu"					,     lowerCAmelCase_=0.1					,     lowerCAmelCase_=0.1					,     lowerCAmelCase_=5_12					,     lowerCAmelCase_=16					,     lowerCAmelCase_=2					,     lowerCAmelCase_=0.02					,     lowerCAmelCase_=3					,     lowerCAmelCase_=4					,     lowerCAmelCase_=None					,     ):





								"""simple docstring"""






								_snake_case     		=	parent
								_snake_case     		=	batch_size
								_snake_case     		=	seq_length
								_snake_case     		=	is_training
								_snake_case     		=	use_token_type_ids
								_snake_case     		=	use_labels
								_snake_case     		=	vocab_size
								_snake_case     		=	hidden_size
								_snake_case     		=	num_hidden_layers
								_snake_case     		=	num_attention_heads
								_snake_case     		=	intermediate_size
								_snake_case     		=	hidden_act
								_snake_case     		=	hidden_dropout_prob
								_snake_case     		=	attention_probs_dropout_prob
								_snake_case     		=	max_position_embeddings
								_snake_case     		=	type_vocab_size
								_snake_case     		=	type_sequence_label_size
								_snake_case     		=	initializer_range
								_snake_case     		=	num_labels
								_snake_case     		=	num_choices
								_snake_case     		=	scope
								_snake_case     		=	self.vocab_size - 1


							def 				lowerCamelCase       ( self						):





								"""simple docstring"""






								_snake_case     		=	ids_tensor([self.batch_size, self.seq_length]					,     self.vocab_size						)

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

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

								_snake_case     		=	OpenAIGPTConfig(
								    vocab_size=self.vocab_size					,     n_embd=self.hidden_size					,     n_layer=self.num_hidden_layers					,     n_head=self.num_attention_heads					,     n_positions=self.max_position_embeddings					,     pad_token_id=self.pad_token_id					,     )

								_snake_case     		=	ids_tensor([self.num_hidden_layers, self.num_attention_heads]					,     2						)

								return (
								    config,
								    input_ids,
								    head_mask,
								    token_type_ids,
								    sequence_labels,
								    token_labels,
								    choice_labels,
								)


							def 				lowerCamelCase       ( self					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     *lowerCAmelCase_						):





								"""simple docstring"""






								_snake_case     		=	OpenAIGPTModel(config=lowerCAmelCase_						)
								model.to(lowerCAmelCase_						)
								model.eval()

								_snake_case     		=	model(lowerCAmelCase_					,     token_type_ids=lowerCAmelCase_					,     head_mask=lowerCAmelCase_						)
								_snake_case     		=	model(lowerCAmelCase_					,     token_type_ids=lowerCAmelCase_						)
								_snake_case     		=	model(lowerCAmelCase_						)

								self.parent.assertEqual(result.last_hidden_state.shape					,     (self.batch_size, self.seq_length, self.hidden_size)						)


							def 				lowerCamelCase       ( self					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     *lowerCAmelCase_						):





								"""simple docstring"""






								_snake_case     		=	OpenAIGPTLMHeadModel(lowerCAmelCase_						)
								model.to(lowerCAmelCase_						)
								model.eval()

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


							def 				lowerCamelCase       ( self					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     *lowerCAmelCase_						):





								"""simple docstring"""






								_snake_case     		=	OpenAIGPTDoubleHeadsModel(lowerCAmelCase_						)
								model.to(lowerCAmelCase_						)
								model.eval()

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


							def 				lowerCamelCase       ( self					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     *lowerCAmelCase_						):





								"""simple docstring"""






								_snake_case     		=	self.num_labels
								_snake_case     		=	OpenAIGPTForSequenceClassification(lowerCAmelCase_						)
								model.to(lowerCAmelCase_						)
								model.eval()

								_snake_case     		=	ids_tensor([self.batch_size]					,     self.type_sequence_label_size						)
								_snake_case     		=	model(lowerCAmelCase_					,     token_type_ids=lowerCAmelCase_					,     labels=lowerCAmelCase_						)
								self.parent.assertEqual(result.logits.shape					,     (self.batch_size, self.num_labels)						)




							def 				lowerCamelCase       ( self						):





								"""simple docstring"""






								_snake_case     		=	self.prepare_config_and_inputs()
								(
								    (
								    _snake_case
								)		,				(
								    _snake_case
								)		,				(
								    _snake_case
								)		,				(
								    _snake_case
								)		,				(
								    _snake_case
								)		,				(
								    _snake_case
								)		,				(
								    _snake_case
								)		,				
								)     		=	config_and_inputs
								_snake_case     		=	{
								    'input_ids': input_ids,
								    'token_type_ids': token_type_ids,
								    'head_mask': head_mask,
								}

								return config, inputs_dict


@require_torch
class 		__UpperCAmelCase   (      _lowerCamelCase		,     _lowerCamelCase		,     _lowerCamelCase		,     unittest.TestCase		):
							__lowercase   	= (
							    (OpenAIGPTModel, OpenAIGPTLMHeadModel, OpenAIGPTDoubleHeadsModel, OpenAIGPTForSequenceClassification)
							    if is_torch_available()
							    else ()
							)
							__lowercase   	= (
							    (OpenAIGPTLMHeadModel,) if is_torch_available() else ()
							)  # TODO (PVP): Add Double HeadsModel when generate() function is changed accordingly
							__lowercase   	= (
							    {
							        """feature-extraction""": OpenAIGPTModel,
							        """text-classification""": OpenAIGPTForSequenceClassification,
							        """text-generation""": OpenAIGPTLMHeadModel,
							        """zero-shot""": OpenAIGPTForSequenceClassification,
							    }
							    if is_torch_available()
							    else {}
							)


							def 				lowerCamelCase       ( self					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_						):





								"""simple docstring"""






								if pipeline_test_casse_name == "ZeroShotClassificationPipelineTests":
									# Get `tokenizer does not have a padding token` error for both fast/slow tokenizers.
									# `OpenAIGPTConfig` was never used in pipeline tests, either because of a missing checkpoint or because a
									# tiny config could not be created.
									return True

								return False


							def 				lowerCamelCase       ( self					,     lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_=False						):





								"""simple docstring"""






								_snake_case     		=	super()._prepare_for_class(lowerCAmelCase_					,     lowerCAmelCase_					,     return_labels=lowerCAmelCase_						)

								if return_labels:
									if model_class.__name__ == "OpenAIGPTDoubleHeadsModel":
										_snake_case     		=	torch.zeros(
										    (self.model_tester.batch_size, self.model_tester.num_choices, self.model_tester.seq_length)					,     dtype=torch.long					,     device=lowerCAmelCase_					,     )
										_snake_case     		=	inputs_dict['labels']
										_snake_case     		=	inputs_dict['labels']
										_snake_case     		=	torch.zeros(
										    (self.model_tester.batch_size, self.model_tester.num_choices)					,     dtype=torch.long					,     device=lowerCAmelCase_					,     )
										_snake_case     		=	torch.zeros(
										    self.model_tester.batch_size					,     dtype=torch.long					,     device=lowerCAmelCase_						)
								return inputs_dict


							def 				lowerCamelCase       ( self						):





								"""simple docstring"""






								_snake_case     		=	OpenAIGPTModelTester(self						)
								_snake_case     		=	ConfigTester(self					,     config_class=lowerCAmelCase_					,     n_embd=37						)


							def 				lowerCamelCase       ( self						):





								"""simple docstring"""






								self.config_tester.run_common_tests()


							def 				lowerCamelCase       ( self						):





								"""simple docstring"""






								_snake_case     		=	self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_openai_gpt_model(*lowerCAmelCase_						)


							def 				lowerCamelCase       ( self						):





								"""simple docstring"""






								_snake_case     		=	self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_lm_head_model(*lowerCAmelCase_						)


							def 				lowerCamelCase       ( self						):





								"""simple docstring"""






								_snake_case     		=	self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_double_lm_head_model(*lowerCAmelCase_						)


							def 				lowerCamelCase       ( self						):





								"""simple docstring"""






								_snake_case     		=	self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_openai_gpt_for_sequence_classification(*lowerCAmelCase_						)




							@slow
							def 				lowerCamelCase       ( self						):





								"""simple docstring"""






								for model_name in OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									_snake_case     		=	OpenAIGPTModel.from_pretrained(lowerCAmelCase_						)
									self.assertIsNotNone(lowerCAmelCase_						)


@require_torch
class 		__UpperCAmelCase   (      unittest.TestCase		):


							@slow
							def 				lowerCamelCase       ( self						):





								"""simple docstring"""






								_snake_case     		=	OpenAIGPTLMHeadModel.from_pretrained('openai-gpt'						)
								model.to(lowerCAmelCase_						)
								_snake_case     		=	torch.tensor([[4_81, 47_35, 5_44]]					,     dtype=torch.long					,     device=lowerCAmelCase_						)  # the president is
								_snake_case     		=	[
								    4_81,
								    47_35,
								    5_44,
								    2_46,
								    9_63,
								    8_70,
								    7_62,
								    2_39,
								    2_44,
								    4_04_77,
								    2_44,
								    2_49,
								    7_19,
								    8_81,
								    4_87,
								    5_44,
								    2_40,
								    2_44,
								    6_03,
								    4_81,
								]  # the president is a very good man. " \n " i\'m sure he is, " said the

								_snake_case     		=	model.generate(lowerCAmelCase_					,     do_sample=lowerCAmelCase_						)
								self.assertListEqual(output_ids[0].tolist()					,     lowerCAmelCase_						)

