"""simple docstring"""





import unittest

from transformers import AutoTokenizer, FalconConfig, 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 (
				    FalconForCausalLM,
				    FalconForQuestionAnswering,
				    FalconForSequenceClassification,
				    FalconForTokenClassification,
				    FalconModel,
				)
class  _snake_case     :






						def __init__(       self      :			Any							,					UpperCAmelCase      :			Union[str, Any]							,					UpperCAmelCase      :			Dict=3							,					UpperCAmelCase      :			str=7							,					UpperCAmelCase      :			Union[str, Any]=True							,					UpperCAmelCase      :			Tuple=True							,					UpperCAmelCase      :			List[str]=False							,					UpperCAmelCase      :			Tuple=True							,					UpperCAmelCase      :			Any=99							,					UpperCAmelCase      :			Any=32							,					UpperCAmelCase      :			str=5							,					UpperCAmelCase      :			List[Any]=4							,					UpperCAmelCase      :			Optional[int]=37							,					UpperCAmelCase      :			str="gelu"							,					UpperCAmelCase      :			Tuple=0.1							,					UpperCAmelCase      :			str=0.1							,					UpperCAmelCase      :			int=512							,					UpperCAmelCase      :			Union[str, Any]=16							,					UpperCAmelCase      :			Union[str, Any]=2							,					UpperCAmelCase      :			Any=0.0_2							,					UpperCAmelCase      :			Any=3							,					UpperCAmelCase      :			str=4							,					UpperCAmelCase      :			Dict=None							,					):
													__lowerCamelCase					:		Union[str, Any]		   =     parent
													__lowerCamelCase					:		List[Any]		   =     batch_size
													__lowerCamelCase					:		str		   =     seq_length
													__lowerCamelCase					:		str		   =     is_training
													__lowerCamelCase					:		Union[str, Any]		   =     use_input_mask
													__lowerCamelCase					:		List[Any]		   =     use_token_type_ids
													__lowerCamelCase					:		Tuple		   =     use_labels
													__lowerCamelCase					:		Any		   =     vocab_size
													__lowerCamelCase					:		Union[str, Any]		   =     hidden_size
													__lowerCamelCase					:		Dict		   =     num_hidden_layers
													__lowerCamelCase					:		Optional[int]		   =     num_attention_heads
													__lowerCamelCase					:		Any		   =     intermediate_size
													__lowerCamelCase					:		str		   =     hidden_act
													__lowerCamelCase					:		Dict		   =     hidden_dropout_prob
													__lowerCamelCase					:		Tuple		   =     attention_probs_dropout_prob
													__lowerCamelCase					:		List[Any]		   =     max_position_embeddings
													__lowerCamelCase					:		Any		   =     type_vocab_size
													__lowerCamelCase					:		List[Any]		   =     type_sequence_label_size
													__lowerCamelCase					:		int		   =     initializer_range
													__lowerCamelCase					:		Tuple		   =     num_labels
													__lowerCamelCase					:		Union[str, Any]		   =     num_choices
													__lowerCamelCase					:		Any		   =     scope






						def       lowerCamelCase__     (       self      :			Any    ):
													__lowerCamelCase					:		Any		   =     ids_tensor([self.batch_size, self.seq_length]							,					self.vocab_size    )

													__lowerCamelCase					:		Dict		   =     None
													if self.use_input_mask:
																				__lowerCamelCase					:		List[str]		   =     random_attention_mask([self.batch_size, self.seq_length]    )

													__lowerCamelCase					:		Dict		   =     None

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

													__lowerCamelCase					:		Tuple		   =     self.get_config()

													return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels






						def       lowerCamelCase__     (       self      :			int    ):
													return FalconConfig(
													    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							,					pad_token_id=1							,					new_decoder_architecture=UpperCAmelCase							,					)






						def       lowerCamelCase__     (       self      :			Tuple							,					UpperCAmelCase      :			Union[str, Any]							,					UpperCAmelCase      :			List[str]							,					UpperCAmelCase      :			str							,					UpperCAmelCase      :			Union[str, Any]							,					UpperCAmelCase      :			str							,					UpperCAmelCase      :			Union[str, Any]							,					UpperCAmelCase      :			Optional[Any]    ):
													__lowerCamelCase					:		Any		   =     FalconModel(config=UpperCAmelCase    )
													model.to(UpperCAmelCase    )
													model.eval()
													__lowerCamelCase					:		Any		   =     model(UpperCAmelCase							,					attention_mask=UpperCAmelCase    )
													__lowerCamelCase					:		List[Any]		   =     model(UpperCAmelCase    )
													self.parent.assertEqual(result.last_hidden_state.shape							,					(self.batch_size, self.seq_length, self.hidden_size)    )






						def       lowerCamelCase__     (       self      :			List[Any]							,					UpperCAmelCase      :			List[str]							,					UpperCAmelCase      :			Tuple							,					UpperCAmelCase      :			str							,					UpperCAmelCase      :			Optional[int]							,					UpperCAmelCase      :			Dict							,					UpperCAmelCase      :			List[str]							,					UpperCAmelCase      :			Any							,					UpperCAmelCase      :			List[Any]							,					UpperCAmelCase      :			List[str]							,					):
													__lowerCamelCase					:		int		   =     True
													__lowerCamelCase					:		Any		   =     FalconModel(UpperCAmelCase    )
													model.to(UpperCAmelCase    )
													model.eval()
													__lowerCamelCase					:		Dict		   =     model(
													    UpperCAmelCase							,					attention_mask=UpperCAmelCase							,					encoder_hidden_states=UpperCAmelCase							,					encoder_attention_mask=UpperCAmelCase							,					)
													__lowerCamelCase					:		Union[str, Any]		   =     model(
													    UpperCAmelCase							,					attention_mask=UpperCAmelCase							,					encoder_hidden_states=UpperCAmelCase							,					)
													__lowerCamelCase					:		List[Any]		   =     model(UpperCAmelCase							,					attention_mask=UpperCAmelCase    )
													self.parent.assertEqual(result.last_hidden_state.shape							,					(self.batch_size, self.seq_length, self.hidden_size)    )






						def       lowerCamelCase__     (       self      :			str							,					UpperCAmelCase      :			Dict							,					UpperCAmelCase      :			Any							,					UpperCAmelCase      :			Optional[int]							,					UpperCAmelCase      :			int							,					UpperCAmelCase      :			Dict							,					UpperCAmelCase      :			Any							,					UpperCAmelCase      :			List[str]							,					UpperCAmelCase      :			List[Any]							,					UpperCAmelCase      :			List[str]							,					):
													__lowerCamelCase					:		Any		   =     FalconForCausalLM(config=UpperCAmelCase    )
													model.to(UpperCAmelCase    )
													model.eval()
													__lowerCamelCase					:		List[Any]		   =     model(UpperCAmelCase							,					attention_mask=UpperCAmelCase							,					labels=UpperCAmelCase    )
													self.parent.assertEqual(result.logits.shape							,					(self.batch_size, self.seq_length, self.vocab_size)    )






						def       lowerCamelCase__     (       self      :			List[Any]							,					UpperCAmelCase      :			Dict							,					UpperCAmelCase      :			List[Any]							,					UpperCAmelCase      :			Any							,					UpperCAmelCase      :			Tuple							,					UpperCAmelCase      :			Dict							,					UpperCAmelCase      :			List[str]							,					UpperCAmelCase      :			str							,					UpperCAmelCase      :			Optional[Any]							,					UpperCAmelCase      :			Union[str, Any]							,					):
													__lowerCamelCase					:		Any		   =     True
													__lowerCamelCase					:		List[Any]		   =     True
													__lowerCamelCase					:		List[str]		   =     FalconForCausalLM(config=UpperCAmelCase    )
													model.to(UpperCAmelCase    )
													model.eval()

													# first forward pass
													__lowerCamelCase					:		Optional[Any]		   =     model(
													    UpperCAmelCase							,					attention_mask=UpperCAmelCase							,					encoder_hidden_states=UpperCAmelCase							,					encoder_attention_mask=UpperCAmelCase							,					use_cache=UpperCAmelCase							,					)
													__lowerCamelCase					:		int		   =     outputs.past_key_values

													# create hypothetical multiple next token and extent to next_input_ids
													__lowerCamelCase					:		List[str]		   =     ids_tensor((self.batch_size, 3)							,					config.vocab_size    )
													__lowerCamelCase					:		str		   =     ids_tensor((self.batch_size, 3)							,					vocab_size=2    )

													# append to next input_ids and
													__lowerCamelCase					:		List[str]		   =     torch.cat([input_ids, next_tokens]							,					dim=-1    )
													__lowerCamelCase					:		Union[str, Any]		   =     torch.cat([input_mask, next_mask]							,					dim=-1    )

													__lowerCamelCase					:		Optional[Any]		   =     model(
													    UpperCAmelCase							,					attention_mask=UpperCAmelCase							,					encoder_hidden_states=UpperCAmelCase							,					encoder_attention_mask=UpperCAmelCase							,					output_hidden_states=UpperCAmelCase							,					)["hidden_states"][0]
													__lowerCamelCase					:		Optional[Any]		   =     model(
													    UpperCAmelCase							,					attention_mask=UpperCAmelCase							,					encoder_hidden_states=UpperCAmelCase							,					encoder_attention_mask=UpperCAmelCase							,					past_key_values=UpperCAmelCase							,					output_hidden_states=UpperCAmelCase							,					)["hidden_states"][0]

													# select random slice
													__lowerCamelCase					:		Any		   =     ids_tensor((1,)							,					output_from_past.shape[-1]    ).item()
													__lowerCamelCase					:		str		   =     output_from_no_past[:, -3:, random_slice_idx].detach()
													__lowerCamelCase					:		Union[str, Any]		   =     output_from_past[:, :, random_slice_idx].detach()

													self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]    )

													# test that outputs are equal for slice
													self.parent.assertTrue(torch.allclose(UpperCAmelCase							,					UpperCAmelCase							,					atol=1E-3    )    )







						def       lowerCamelCase__     (       self      :			Optional[int]    ):
													__lowerCamelCase					:		Dict		   =     self.prepare_config_and_inputs()
													(
													    (
													    __lowerCamelCase
													)   ,  (
													    __lowerCamelCase
													)   ,  (
													    __lowerCamelCase
													)   ,  (
													    __lowerCamelCase
													)   ,  (
													    __lowerCamelCase
													)   ,  (
													    __lowerCamelCase
													)   ,  (
													    __lowerCamelCase
													)   ,  
													)					:		str		   =     config_and_inputs
													__lowerCamelCase					:		Dict		   =     {"input_ids": input_ids, "attention_mask": input_mask}
													return config, inputs_dict

@require_torch
class  _snake_case     (					a__				, a__				, a__				, unittest.TestCase		):
						snake_case__			  =			(
						    (
						        FalconModel,
						        FalconForCausalLM,
						        FalconForSequenceClassification,
						        FalconForTokenClassification,
						        FalconForQuestionAnswering,
						    )
						    if is_torch_available()
						    else ()
						)
						snake_case__			  =			(FalconForCausalLM,) if is_torch_available() else ()
						snake_case__			  =			(
						    {
						        "feature-extraction": FalconModel,
						        "text-classification": FalconForSequenceClassification,
						        "text-generation": FalconForCausalLM,
						        "question-answering": FalconForQuestionAnswering,
						        "token-classification": FalconForTokenClassification,
						        "zero-shot": FalconForSequenceClassification,
						    }
						    if is_torch_available()
						    else {}
						)
						snake_case__			  =			False
						snake_case__			  =			False






						def       lowerCamelCase__     (       self      :			List[Any]    ):
													__lowerCamelCase					:		Dict		   =     FalconModelTester(self    )
													__lowerCamelCase					:		Dict		   =     ConfigTester(self							,					config_class=UpperCAmelCase							,					hidden_size=37    )






						def       lowerCamelCase__     (       self      :			List[str]    ):
													self.config_tester.run_common_tests()






						def       lowerCamelCase__     (       self      :			List[Any]    ):
													__lowerCamelCase					:		Dict		   =     self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_model(*UpperCAmelCase    )






						def       lowerCamelCase__     (       self      :			Union[str, Any]    ):
													__lowerCamelCase   ,  *__lowerCamelCase					:		str		   =     self.model_tester.prepare_config_and_inputs()
													for alibi in [True, False]:
																				__lowerCamelCase					:		str		   =     alibi
																				self.model_tester.create_and_check_model(UpperCAmelCase							,					*UpperCAmelCase    )






						def       lowerCamelCase__     (       self      :			str    ):
													__lowerCamelCase   ,  __lowerCamelCase					:		Any		   =     self.model_tester.prepare_config_and_inputs_for_common()
													__lowerCamelCase					:		Union[str, Any]		   =     3
													__lowerCamelCase					:		Union[str, Any]		   =     input_dict["input_ids"]
													__lowerCamelCase					:		int		   =     input_ids.ne(1    ).to(UpperCAmelCase    )
													__lowerCamelCase					:		int		   =     ids_tensor([self.model_tester.batch_size]							,					self.model_tester.type_sequence_label_size    )
													__lowerCamelCase					:		Optional[int]		   =     FalconForSequenceClassification(UpperCAmelCase    )
													model.to(UpperCAmelCase    )
													model.eval()
													__lowerCamelCase					:		Union[str, Any]		   =     model(UpperCAmelCase							,					attention_mask=UpperCAmelCase							,					labels=UpperCAmelCase    )
													self.assertEqual(result.logits.shape							,					(self.model_tester.batch_size, self.model_tester.num_labels)    )






						def       lowerCamelCase__     (       self      :			Union[str, Any]    ):
													__lowerCamelCase   ,  __lowerCamelCase					:		Dict		   =     self.model_tester.prepare_config_and_inputs_for_common()
													__lowerCamelCase					:		List[str]		   =     3
													__lowerCamelCase					:		Optional[int]		   =     "single_label_classification"
													__lowerCamelCase					:		Any		   =     input_dict["input_ids"]
													__lowerCamelCase					:		Optional[int]		   =     input_ids.ne(1    ).to(UpperCAmelCase    )
													__lowerCamelCase					:		Any		   =     ids_tensor([self.model_tester.batch_size]							,					self.model_tester.type_sequence_label_size    )
													__lowerCamelCase					:		Tuple		   =     FalconForSequenceClassification(UpperCAmelCase    )
													model.to(UpperCAmelCase    )
													model.eval()
													__lowerCamelCase					:		Dict		   =     model(UpperCAmelCase							,					attention_mask=UpperCAmelCase							,					labels=UpperCAmelCase    )
													self.assertEqual(result.logits.shape							,					(self.model_tester.batch_size, self.model_tester.num_labels)    )






						def       lowerCamelCase__     (       self      :			Optional[int]    ):
													__lowerCamelCase   ,  __lowerCamelCase					:		Dict		   =     self.model_tester.prepare_config_and_inputs_for_common()
													__lowerCamelCase					:		Tuple		   =     input_dict["input_ids"]
													__lowerCamelCase					:		str		   =     FalconForCausalLM(UpperCAmelCase    )
													model.to(UpperCAmelCase    )
													model.eval()
													__lowerCamelCase					:		int		   =     model(UpperCAmelCase							,					use_cache=UpperCAmelCase    )
													__lowerCamelCase					:		Union[str, Any]		   =     input_ids.shape[0]
													__lowerCamelCase					:		int		   =     model._convert_to_rw_cache(result.past_key_values    )
													__lowerCamelCase					:		Union[str, Any]		   =     model._convert_cache_to_standard_format(UpperCAmelCase							,					UpperCAmelCase    )
													for layer in range(len(UpperCAmelCase    )    ):
																				for tensor_idx in range(2    ):
																											self.assertTrue(rw_cache[layer][tensor_idx].ndim == 3    )
																											self.assertTrue(result.past_key_values[layer][tensor_idx].ndim == 4    )
																											self.assertTrue(
																											    torch.all(result.past_key_values[layer][tensor_idx] == standard_cache[layer][tensor_idx]    )    )






						def       lowerCamelCase__     (       self      :			int    ):
													__lowerCamelCase   ,  __lowerCamelCase					:		Optional[Any]		   =     self.model_tester.prepare_config_and_inputs_for_common()
													__lowerCamelCase					:		List[Any]		   =     3
													__lowerCamelCase					:		Tuple		   =     "multi_label_classification"
													__lowerCamelCase					:		List[Any]		   =     input_dict["input_ids"]
													__lowerCamelCase					:		Tuple		   =     input_ids.ne(1    ).to(UpperCAmelCase    )
													__lowerCamelCase					:		List[Any]		   =     ids_tensor(
													    [self.model_tester.batch_size, config.num_labels]							,					self.model_tester.type_sequence_label_size    ).to(torch.float    )
													__lowerCamelCase					:		Union[str, Any]		   =     FalconForSequenceClassification(UpperCAmelCase    )
													model.to(UpperCAmelCase    )
													model.eval()
													__lowerCamelCase					:		Any		   =     model(UpperCAmelCase							,					attention_mask=UpperCAmelCase							,					labels=UpperCAmelCase    )
													self.assertEqual(result.logits.shape							,					(self.model_tester.batch_size, self.model_tester.num_labels)    )







						def       lowerCamelCase__     (       self      :			List[str]    ):
													# Falcon can have different numbers of KV-heads than the number of query heads, so we need
													# to override this test to use the right head counts.
													for model_class in self.all_generative_model_classes:
																				__lowerCamelCase   ,  __lowerCamelCase					:		Optional[int]		   =     self.model_tester.prepare_config_and_inputs_for_common()

																				# If it doesn't support cache, pass the test
																				if not hasattr(UpperCAmelCase							,					"use_cache"    ):
																											return

																				__lowerCamelCase					:		Any		   =     model_class(UpperCAmelCase    ).to(UpperCAmelCase    )
																				if "use_cache" not in inputs:
																											__lowerCamelCase					:		List[str]		   =     True
																				__lowerCamelCase					:		int		   =     model(**UpperCAmelCase    )

																				# If "past_key_values" is not returned, pass the test (e.g. RWKV uses a different cache name and format)
																				if "past_key_values" not in outputs:
																											return

																				__lowerCamelCase					:		List[Any]		   =     (
																				    getattr(UpperCAmelCase							,					"decoder_layers"							,					UpperCAmelCase    )
																				    or getattr(UpperCAmelCase							,					"num_decoder_layers"							,					UpperCAmelCase    )
																				    or config.num_hidden_layers
																				)
																				__lowerCamelCase					:		Union[str, Any]		   =     getattr(UpperCAmelCase							,					"num_kv_heads"							,					config.num_attention_heads    )
																				__lowerCamelCase					:		List[Any]		   =     getattr(UpperCAmelCase							,					"d_model"							,					config.hidden_size    )
																				__lowerCamelCase					:		Any		   =     embed_dim // num_attention_heads

																				__lowerCamelCase					:		List[Any]		   =     outputs["past_key_values"]
																				self.assertEqual(len(UpperCAmelCase    )							,					UpperCAmelCase    )

																				__lowerCamelCase   ,  __lowerCamelCase					:		List[Any]		   =     inputs["input_ids"].shape
																				for i in range(UpperCAmelCase    ):
																											if config.new_decoder_architecture:
																																		__lowerCamelCase					:		int		   =     config.num_attention_heads
																											elif config.multi_query:
																																		__lowerCamelCase					:		Optional[int]		   =     1
																											self.assertEqual(len(past_kv[0]    )							,					2    )  # K V for the decoder = 2
																											self.assertEqual(
																											    past_kv[i][0].shape							,					(batch_size, num_attention_heads, seq_length, per_head_embed_dim)    )
																											self.assertEqual(
																											    past_kv[i][1].shape							,					(batch_size, num_attention_heads, seq_length, per_head_embed_dim)    )

@require_torch
class  _snake_case     (					unittest.TestCase		):






						@slow
						def       lowerCamelCase__     (       self      :			Any    ):
													__lowerCamelCase					:		Optional[Any]		   =     AutoTokenizer.from_pretrained("Rocketknight1/falcon-rw-1b"    )
													__lowerCamelCase					:		List[str]		   =     FalconForCausalLM.from_pretrained("Rocketknight1/falcon-rw-1b"    )
													model.eval()
													model.to(UpperCAmelCase    )
													__lowerCamelCase					:		str		   =     tokenizer("My favorite food is"							,					return_tensors="pt"    ).to(UpperCAmelCase    )

													__lowerCamelCase					:		Optional[int]		   =     (
													    "My favorite food is pizza. I love it so much that I have a pizza party every year for my birthday."
													)

													__lowerCamelCase					:		Union[str, Any]		   =     model.generate(**UpperCAmelCase							,					do_sample=UpperCAmelCase							,					max_new_tokens=19    )
													__lowerCamelCase					:		Union[str, Any]		   =     tokenizer.batch_decode(UpperCAmelCase    )[0]

													self.assertEqual(UpperCAmelCase							,					UpperCAmelCase    )






						@slow
						def       lowerCamelCase__     (       self      :			List[str]    ):
													# The big models are way too big for the CI, so we use tiny random models that resemble their
													# architectures but with much smaller and fewer layers
													for repo in ["Rocketknight1/tiny-random-falcon-7b", "Rocketknight1/tiny-random-falcon-40b"]:
																				__lowerCamelCase					:		List[Any]		   =     AutoTokenizer.from_pretrained(UpperCAmelCase    )
																				__lowerCamelCase					:		List[Any]		   =     FalconForCausalLM.from_pretrained(UpperCAmelCase    )
																				model.eval()
																				model.to(UpperCAmelCase    )
																				__lowerCamelCase					:		int		   =     tokenizer("My favorite food is"							,					return_tensors="pt"    ).to(UpperCAmelCase    )

																				# We just test that these run without errors - the models are randomly initialized
																				# and so the actual text outputs will be garbage
																				model.generate(**UpperCAmelCase							,					do_sample=UpperCAmelCase							,					max_new_tokens=4    )
																				model.generate(**UpperCAmelCase							,					do_sample=UpperCAmelCase							,					max_new_tokens=4    )
																				model.generate(**UpperCAmelCase							,					num_beams=2							,					max_new_tokens=4    )







						@slow
						def       lowerCamelCase__     (       self      :			List[str]    ):
													# The big models are way too big for the CI, so we use tiny random models that resemble their
													# architectures but with much smaller and fewer layers
													with torch.no_grad():
																				for repo in [
																				    "Rocketknight1/falcon-rw-1b",
																				    "Rocketknight1/tiny-random-falcon-7b",
																				    "Rocketknight1/tiny-random-falcon-40b",
																				]:
																											__lowerCamelCase					:		Tuple		   =     AutoTokenizer.from_pretrained(UpperCAmelCase    )
																											__lowerCamelCase					:		List[Any]		   =     FalconForCausalLM.from_pretrained(UpperCAmelCase    )
																											model.eval()
																											model.to(device=UpperCAmelCase    )
																											__lowerCamelCase					:		List[Any]		   =     tokenizer("My favorite food is"							,					return_tensors="pt"    ).to(UpperCAmelCase    )

																											# Test results are the same with and without cache
																											__lowerCamelCase					:		List[Any]		   =     model.generate(**UpperCAmelCase							,					do_sample=UpperCAmelCase							,					max_new_tokens=20							,					use_cache=UpperCAmelCase    )
																											__lowerCamelCase					:		Optional[int]		   =     model.generate(**UpperCAmelCase							,					do_sample=UpperCAmelCase							,					max_new_tokens=20							,					use_cache=UpperCAmelCase    )
																											self.assertTrue((outputs_cache - outputs_no_cache).sum().item() == 0    )
"""simple docstring"""





import gc
import threading
import time

import psutil
import torch
class  _snake_case     :






						def __init__(       self      :			str    ):
													__lowerCamelCase					:		Optional[Any]		   =     psutil.Process()
													__lowerCamelCase					:		List[Any]		   =     False






						def       lowerCamelCase__     (       self      :			str    ):
													__lowerCamelCase					:		List[Any]		   =     -1

													while True:
																				__lowerCamelCase					:		Union[str, Any]		   =     max(self.process.memory_info().rss							,					self.cpu_memory_peak    )

																				# can't sleep or will not catch the peak right (this comment is here on purpose)
																				if not self.peak_monitoring:
																											break






						def       lowerCamelCase__     (       self      :			Any    ):
													__lowerCamelCase					:		Optional[Any]		   =     True
													__lowerCamelCase					:		Union[str, Any]		   =     threading.Thread(target=self.peak_monitor    )
													__lowerCamelCase					:		Optional[Any]		   =     True
													self.thread.start()







						def       lowerCamelCase__     (       self      :			Dict    ):
													__lowerCamelCase					:		str		   =     False
													self.thread.join()
													return self.cpu_memory_peak

__A  =				PeakCPUMemory()


def    lowercase_   (   ) ->      Union[str, Any]:



							'''simple docstring'''

							__lowerCamelCase					:		Any		   =     {"time": time.time()}

							gc.collect()
							torch.cuda.empty_cache()

							# CPU mem
							__lowerCamelCase					:		Optional[int]		   =     psutil.Process().memory_info().rss
							cpu_peak_tracker.start()

							# GPU mem
							for i in range(torch.cuda.device_count()      ):
														__lowerCamelCase					:		int		   =     torch.cuda.memory_allocated(_lowerCamelCase      )
							torch.cuda.reset_peak_memory_stats()

							return measures


def    lowercase_   (   _lowerCamelCase:	Optional[Any]      ) ->      Tuple:



							'''simple docstring'''

							__lowerCamelCase					:		Tuple		   =     {"time": time.time() - start_measures["time"]}

							gc.collect()
							torch.cuda.empty_cache()

							# CPU mem
							__lowerCamelCase					:		Dict		   =     (psutil.Process().memory_info().rss - start_measures["cpu"]) / 2**20
							__lowerCamelCase					:		Tuple		   =     (cpu_peak_tracker.stop() - start_measures["cpu"]) / 2**20

							# GPU mem
							for i in range(torch.cuda.device_count()      ):
														__lowerCamelCase					:		Any		   =     (torch.cuda.memory_allocated(_lowerCamelCase      ) - start_measures[str(_lowerCamelCase      )]) / 2**20
														__lowerCamelCase					:		Optional[int]		   =     (torch.cuda.max_memory_allocated(_lowerCamelCase      ) - start_measures[str(_lowerCamelCase      )]) / 2**20

							return measures







def    lowercase_   (   _lowerCamelCase:	int						,   _lowerCamelCase:	List[Any]      ) ->      Optional[int]:



							'''simple docstring'''

							print(F"""{description}:"""      )
							print(F"""- Time: {measures["time"]:.2f}s"""      )
							for i in range(torch.cuda.device_count()      ):
														print(F"""- GPU {i} allocated: {measures[str(_lowerCamelCase      )]:.2f}MiB"""      )
														__lowerCamelCase					:		List[Any]		   =     measures[F"""{i}-peak"""]
														print(F"""- GPU {i} peak: {peak:.2f}MiB"""      )
							print(F"""- CPU RAM allocated: {measures["cpu"]:.2f}MiB"""      )
							print(F"""- CPU RAM peak: {measures["cpu-peak"]:.2f}MiB"""      )