
'''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 						a   :




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




				def 	__snake_case							(							self	:				List[Any]      ):
										snake_case       :     int =    ids_tensor([self.batch_size, self.seq_length],    self.vocab_size      )

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

										snake_case       :     List[str] =    None

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

										snake_case       :     Optional[Any] =    self.get_config()

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




				def 	__snake_case							(							self	:				Optional[Any]      ):
										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=SCREAMING_SNAKE_CASE_,    initializer_range=self.initializer_range,    pad_token_id=1,    new_decoder_architecture=SCREAMING_SNAKE_CASE_,    )




				def 	__snake_case							(							self	:				List[str],    SCREAMING_SNAKE_CASE_	:				List[Any],    SCREAMING_SNAKE_CASE_	:				Tuple,    SCREAMING_SNAKE_CASE_	:				Optional[Any],    SCREAMING_SNAKE_CASE_	:				Union[str, Any],    SCREAMING_SNAKE_CASE_	:				List[str],    SCREAMING_SNAKE_CASE_	:				int,    SCREAMING_SNAKE_CASE_	:				List[str]      ):
										snake_case       :     Union[str, Any] =    FalconModel(config=SCREAMING_SNAKE_CASE_      )
										model.to(SCREAMING_SNAKE_CASE_      )
										model.eval()
										snake_case       :     Dict =    model(SCREAMING_SNAKE_CASE_,    attention_mask=SCREAMING_SNAKE_CASE_      )
										snake_case       :     Dict =    model(SCREAMING_SNAKE_CASE_      )
										self.parent.assertEqual(result.last_hidden_state.shape,    (self.batch_size, self.seq_length, self.hidden_size)      )




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




				def 	__snake_case							(							self	:				List[str],    SCREAMING_SNAKE_CASE_	:				List[Any],    SCREAMING_SNAKE_CASE_	:				str,    SCREAMING_SNAKE_CASE_	:				Dict,    SCREAMING_SNAKE_CASE_	:				Union[str, Any],    SCREAMING_SNAKE_CASE_	:				Optional[int],    SCREAMING_SNAKE_CASE_	:				Any,    SCREAMING_SNAKE_CASE_	:				Optional[Any],    SCREAMING_SNAKE_CASE_	:				Any,    SCREAMING_SNAKE_CASE_	:				str,    ):
										snake_case       :     Tuple =    FalconForCausalLM(config=SCREAMING_SNAKE_CASE_      )
										model.to(SCREAMING_SNAKE_CASE_      )
										model.eval()
										snake_case       :     int =    model(SCREAMING_SNAKE_CASE_,    attention_mask=SCREAMING_SNAKE_CASE_,    labels=SCREAMING_SNAKE_CASE_      )
										self.parent.assertEqual(result.logits.shape,    (self.batch_size, self.seq_length, self.vocab_size)      )




				def 	__snake_case							(							self	:				Union[str, Any],    SCREAMING_SNAKE_CASE_	:				Dict,    SCREAMING_SNAKE_CASE_	:				Optional[int],    SCREAMING_SNAKE_CASE_	:				Optional[Any],    SCREAMING_SNAKE_CASE_	:				int,    SCREAMING_SNAKE_CASE_	:				Optional[int],    SCREAMING_SNAKE_CASE_	:				Any,    SCREAMING_SNAKE_CASE_	:				List[str],    SCREAMING_SNAKE_CASE_	:				Optional[int],    SCREAMING_SNAKE_CASE_	:				List[Any],    ):
										snake_case       :     Optional[int] =    True
										snake_case       :     str =    True
										snake_case       :     int =    FalconForCausalLM(config=SCREAMING_SNAKE_CASE_      )
										model.to(SCREAMING_SNAKE_CASE_      )
										model.eval()

										# first forward pass
										snake_case       :     Tuple =    model(
										    SCREAMING_SNAKE_CASE_,    attention_mask=SCREAMING_SNAKE_CASE_,    encoder_hidden_states=SCREAMING_SNAKE_CASE_,    encoder_attention_mask=SCREAMING_SNAKE_CASE_,    use_cache=SCREAMING_SNAKE_CASE_,    )
										snake_case       :     Tuple =    outputs.past_key_values

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

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

										snake_case       :     Union[str, Any] =    model(
										    SCREAMING_SNAKE_CASE_,    attention_mask=SCREAMING_SNAKE_CASE_,    encoder_hidden_states=SCREAMING_SNAKE_CASE_,    encoder_attention_mask=SCREAMING_SNAKE_CASE_,    output_hidden_states=SCREAMING_SNAKE_CASE_,    )['''hidden_states'''][0]
										snake_case       :     Tuple =    model(
										    SCREAMING_SNAKE_CASE_,    attention_mask=SCREAMING_SNAKE_CASE_,    encoder_hidden_states=SCREAMING_SNAKE_CASE_,    encoder_attention_mask=SCREAMING_SNAKE_CASE_,    past_key_values=SCREAMING_SNAKE_CASE_,    output_hidden_states=SCREAMING_SNAKE_CASE_,    )['''hidden_states'''][0]

										# select random slice
										snake_case       :     List[str] =    ids_tensor((1,),    output_from_past.shape[-1]      ).item()
										snake_case       :     Union[str, Any] =    output_from_no_past[:, -3:, random_slice_idx].detach()
										snake_case       :     Optional[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(SCREAMING_SNAKE_CASE_,    SCREAMING_SNAKE_CASE_,    atol=1e-3      )      )




				def 	__snake_case							(							self	:				int      ):
										snake_case       :     Any =    self.prepare_config_and_inputs()
										(
										    (
										    snake_case
										),					(
										    snake_case
										),					(
										    snake_case
										),					(
										    snake_case
										),					(
										    snake_case
										),					(
										    snake_case
										),					(
										    snake_case
										),					
										)       :     str =    config_and_inputs
										snake_case       :     Union[str, Any] =    {'''input_ids''': input_ids, '''attention_mask''': input_mask}
										return config, inputs_dict




@require_torch
class 						a   (				__magic_name__				,__magic_name__				,__magic_name__				,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 	__snake_case							(							self	:				List[str]      ):
										snake_case       :     str =    FalconModelTester(self      )
										snake_case       :     Optional[int] =    ConfigTester(self,    config_class=SCREAMING_SNAKE_CASE_,    hidden_size=37      )




				def 	__snake_case							(							self	:				Optional[int]      ):
										self.config_tester.run_common_tests()




				def 	__snake_case							(							self	:				Optional[Any]      ):
										snake_case       :     Optional[Any] =    self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE_      )




				def 	__snake_case							(							self	:				Optional[Any]      ):
										snake_case,					*snake_case       :     Optional[Any] =    self.model_tester.prepare_config_and_inputs()
										for alibi in [True, False]:
																snake_case       :     Dict =    alibi
																self.model_tester.create_and_check_model(SCREAMING_SNAKE_CASE_,    *SCREAMING_SNAKE_CASE_      )




				def 	__snake_case							(							self	:				int      ):
										snake_case,					snake_case       :     List[Any] =    self.model_tester.prepare_config_and_inputs_for_common()
										snake_case       :     Dict =    3
										snake_case       :     Any =    input_dict['''input_ids''']
										snake_case       :     Tuple =    input_ids.ne(1      ).to(SCREAMING_SNAKE_CASE_      )
										snake_case       :     Optional[int] =    ids_tensor([self.model_tester.batch_size],    self.model_tester.type_sequence_label_size      )
										snake_case       :     str =    FalconForSequenceClassification(SCREAMING_SNAKE_CASE_      )
										model.to(SCREAMING_SNAKE_CASE_      )
										model.eval()
										snake_case       :     Optional[int] =    model(SCREAMING_SNAKE_CASE_,    attention_mask=SCREAMING_SNAKE_CASE_,    labels=SCREAMING_SNAKE_CASE_      )
										self.assertEqual(result.logits.shape,    (self.model_tester.batch_size, self.model_tester.num_labels)      )




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




				def 	__snake_case							(							self	:				int      ):
										snake_case,					snake_case       :     List[Any] =    self.model_tester.prepare_config_and_inputs_for_common()
										snake_case       :     Tuple =    input_dict['''input_ids''']
										snake_case       :     int =    FalconForCausalLM(SCREAMING_SNAKE_CASE_      )
										model.to(SCREAMING_SNAKE_CASE_      )
										model.eval()
										snake_case       :     List[Any] =    model(SCREAMING_SNAKE_CASE_,    use_cache=SCREAMING_SNAKE_CASE_      )
										snake_case       :     Any =    input_ids.shape[0]
										snake_case       :     Dict =    model._convert_to_rw_cache(result.past_key_values      )
										snake_case       :     int =    model._convert_cache_to_standard_format(SCREAMING_SNAKE_CASE_,    SCREAMING_SNAKE_CASE_      )
										for layer in range(len(SCREAMING_SNAKE_CASE_      )      ):
																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 	__snake_case							(							self	:				str      ):
										snake_case,					snake_case       :     Dict =    self.model_tester.prepare_config_and_inputs_for_common()
										snake_case       :     Tuple =    3
										snake_case       :     int =    '''multi_label_classification'''
										snake_case       :     Union[str, Any] =    input_dict['''input_ids''']
										snake_case       :     Union[str, Any] =    input_ids.ne(1      ).to(SCREAMING_SNAKE_CASE_      )
										snake_case       :     Optional[int] =    ids_tensor(
										    [self.model_tester.batch_size, config.num_labels],    self.model_tester.type_sequence_label_size      ).to(torch.float      )
										snake_case       :     Union[str, Any] =    FalconForSequenceClassification(SCREAMING_SNAKE_CASE_      )
										model.to(SCREAMING_SNAKE_CASE_      )
										model.eval()
										snake_case       :     List[Any] =    model(SCREAMING_SNAKE_CASE_,    attention_mask=SCREAMING_SNAKE_CASE_,    labels=SCREAMING_SNAKE_CASE_      )
										self.assertEqual(result.logits.shape,    (self.model_tester.batch_size, self.model_tester.num_labels)      )




				def 	__snake_case							(							self	:				Optional[Any]      ):
										# 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:
																snake_case,					snake_case       :     Optional[int] =    self.model_tester.prepare_config_and_inputs_for_common()

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

																snake_case       :     Union[str, Any] =    model_class(SCREAMING_SNAKE_CASE_      ).to(SCREAMING_SNAKE_CASE_      )
																if "use_cache" not in inputs:
																						snake_case       :     Union[str, Any] =    True
																snake_case       :     List[Any] =    model(**SCREAMING_SNAKE_CASE_      )

																# 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

																snake_case       :     Tuple =    (
																    getattr(SCREAMING_SNAKE_CASE_,    '''decoder_layers''',    SCREAMING_SNAKE_CASE_      )
																    or getattr(SCREAMING_SNAKE_CASE_,    '''num_decoder_layers''',    SCREAMING_SNAKE_CASE_      )
																    or config.num_hidden_layers
																)
																snake_case       :     Union[str, Any] =    getattr(SCREAMING_SNAKE_CASE_,    '''num_kv_heads''',    config.num_attention_heads      )
																snake_case       :     Any =    getattr(SCREAMING_SNAKE_CASE_,    '''d_model''',    config.hidden_size      )
																snake_case       :     Any =    embed_dim // num_attention_heads

																snake_case       :     Any =    outputs['''past_key_values''']
																self.assertEqual(len(SCREAMING_SNAKE_CASE_      ),    SCREAMING_SNAKE_CASE_      )

																snake_case,					snake_case       :     Optional[Any] =    inputs['''input_ids'''].shape
																for i in range(SCREAMING_SNAKE_CASE_      ):
																						if config.new_decoder_architecture:
																												snake_case       :     Union[str, Any] =    config.num_attention_heads
																						elif config.multi_query:
																												snake_case       :     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 						a   (				unittest.TestCase    ):




				@slow
				def 	__snake_case							(							self	:				Tuple      ):
										snake_case       :     Optional[Any] =    AutoTokenizer.from_pretrained('''Rocketknight1/falcon-rw-1b'''      )
										snake_case       :     Tuple =    FalconForCausalLM.from_pretrained('''Rocketknight1/falcon-rw-1b'''      )
										model.eval()
										model.to(SCREAMING_SNAKE_CASE_      )
										snake_case       :     Any =    tokenizer('''My favorite food is''',    return_tensors='''pt'''      ).to(SCREAMING_SNAKE_CASE_      )

										snake_case       :     Union[str, Any] =    (
										    '''My favorite food is pizza. I love it so much that I have a pizza party every year for my birthday.'''
										)

										snake_case       :     Optional[Any] =    model.generate(**SCREAMING_SNAKE_CASE_,    do_sample=SCREAMING_SNAKE_CASE_,    max_new_tokens=19      )
										snake_case       :     List[str] =    tokenizer.batch_decode(SCREAMING_SNAKE_CASE_      )[0]

										self.assertEqual(SCREAMING_SNAKE_CASE_,    SCREAMING_SNAKE_CASE_      )




				@slow
				def 	__snake_case							(							self	:				Optional[Any]      ):
										# 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"]:
																snake_case       :     Dict =    AutoTokenizer.from_pretrained(SCREAMING_SNAKE_CASE_      )
																snake_case       :     Any =    FalconForCausalLM.from_pretrained(SCREAMING_SNAKE_CASE_      )
																model.eval()
																model.to(SCREAMING_SNAKE_CASE_      )
																snake_case       :     Dict =    tokenizer('''My favorite food is''',    return_tensors='''pt'''      ).to(SCREAMING_SNAKE_CASE_      )

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




				@slow
				def 	__snake_case							(							self	:				Optional[int]      ):
										# 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",
																]:
																						snake_case       :     Optional[Any] =    AutoTokenizer.from_pretrained(SCREAMING_SNAKE_CASE_      )
																						snake_case       :     Tuple =    FalconForCausalLM.from_pretrained(SCREAMING_SNAKE_CASE_      )
																						model.eval()
																						model.to(device=SCREAMING_SNAKE_CASE_      )
																						snake_case       :     List[str] =    tokenizer('''My favorite food is''',    return_tensors='''pt'''      ).to(SCREAMING_SNAKE_CASE_      )

																						# Test results are the same with and without cache
																						snake_case       :     Union[str, Any] =    model.generate(**SCREAMING_SNAKE_CASE_,    do_sample=SCREAMING_SNAKE_CASE_,    max_new_tokens=20,    use_cache=SCREAMING_SNAKE_CASE_      )
																						snake_case       :     Tuple =    model.generate(**SCREAMING_SNAKE_CASE_,    do_sample=SCREAMING_SNAKE_CASE_,    max_new_tokens=20,    use_cache=SCREAMING_SNAKE_CASE_      )
																						self.assertTrue((outputs_cache - outputs_no_cache).sum().item() == 0      )



'''simple docstring'''





import re
import time
from typing import Optional

import IPython.display as disp

from ..trainer_callback import TrainerCallback
from ..trainer_utils import IntervalStrategy, has_length

def      A (       A_							:					str       ):
						snake_case       :     List[str] =    int(A_       )
						snake_case,					snake_case,					snake_case       :     Any =    t // 3600, (t // 60) % 60, t % 60
						return F"""{h}:{m:02d}:{s:02d}""" if h != 0 else F"""{m:02d}:{s:02d}"""

def      A (       A_							:					Optional[int]    ,							A_							:					int    ,							A_							:					List[str]    ,							A_							:					Optional[int]    ,							A_							:					Optional[int]=300       ):
						# docstyle-ignore
						return F"""
    <div>
      {prefix}
      <progress value='{value}' max='{total}' style='width:{width}px; height:20px; vertical-align: middle;'></progress>
      {label}
    </div>
    """

def      A (       A_							:					Dict       ):
						snake_case       :     Tuple =    '''<table border="1" class="dataframe">\n'''
						html_code += """  <thead>\n <tr style="text-align: left;">\n"""
						for i in items[0]:
												html_code += F"""      <th>{i}</th>\n"""
						html_code += "    </tr>\n  </thead>\n  <tbody>\n"
						for line in items[1:]:
												html_code += "    <tr>\n"
												for elt in line:
																		snake_case       :     List[str] =    F"""{elt:.6f}""" if isinstance(A_    ,							A_       ) else str(A_       )
																		html_code += F"""      <td>{elt}</td>\n"""
												html_code += "    </tr>\n"
						html_code += "  </tbody>\n</table><p>"
						return html_code

class 						a   :
				_snake_case		    = 5
				_snake_case		    = 0.2




				def __init__(							self	:				Union[str, Any],    SCREAMING_SNAKE_CASE_	:				int,    SCREAMING_SNAKE_CASE_	:				Optional[str] = None,    SCREAMING_SNAKE_CASE_	:				bool = True,    SCREAMING_SNAKE_CASE_	:				Optional["NotebookTrainingTracker"] = None,    SCREAMING_SNAKE_CASE_	:				int = 3_00,    ):
										snake_case       :     List[Any] =    total
										snake_case       :     Union[str, Any] =    '''''' if prefix is None else prefix
										snake_case       :     List[Any] =    leave
										snake_case       :     int =    parent
										snake_case       :     List[str] =    width
										snake_case       :     Optional[int] =    None
										snake_case       :     Optional[Any] =    None
										snake_case       :     Tuple =    None




				def 	__snake_case							(							self	:				Any,    SCREAMING_SNAKE_CASE_	:				int,    SCREAMING_SNAKE_CASE_	:				bool = False,    SCREAMING_SNAKE_CASE_	:				str = None      ):
										snake_case       :     List[Any] =    value
										if comment is not None:
																snake_case       :     Tuple =    comment
										if self.last_value is None:
																snake_case       :     str =    time.time()
																snake_case       :     List[str] =    value
																snake_case       :     str =    None
																snake_case       :     Dict =    self.warmup
																snake_case       :     Tuple =    1
																self.update_bar(SCREAMING_SNAKE_CASE_      )
										elif value <= self.last_value and not force_update:
																return
										elif force_update or self.first_calls > 0 or value >= min(self.last_value + self.wait_for,    self.total      ):
																if self.first_calls > 0:
																						self.first_calls -= 1
																snake_case       :     Tuple =    time.time()
																snake_case       :     str =    current_time - self.start_time
																# We could have value = self.start_value if the update is called twixe with the same start value.
																if value > self.start_value:
																						snake_case       :     Any =    self.elapsed_time / (value - self.start_value)
																else:
																						snake_case       :     List[str] =    None
																if value >= self.total:
																						snake_case       :     List[Any] =    self.total
																						snake_case       :     Tuple =    None
																						if not self.leave:
																												self.close()
																elif self.average_time_per_item is not None:
																						snake_case       :     Optional[int] =    self.average_time_per_item * (self.total - value)
																self.update_bar(SCREAMING_SNAKE_CASE_      )
																snake_case       :     List[str] =    value
																snake_case       :     List[Any] =    current_time
																if self.average_time_per_item is None:
																						snake_case       :     List[str] =    1
																else:
																						snake_case       :     Optional[int] =    max(int(self.update_every / self.average_time_per_item      ),    1      )




				def 	__snake_case							(							self	:				Any,    SCREAMING_SNAKE_CASE_	:				List[str],    SCREAMING_SNAKE_CASE_	:				Optional[Any]=None      ):
										snake_case       :     Optional[Any] =    ''' ''' * (len(str(self.total      )      ) - len(str(SCREAMING_SNAKE_CASE_      )      )) + str(SCREAMING_SNAKE_CASE_      )
										if self.elapsed_time is None:
																snake_case       :     int =    F"""[{spaced_value}/{self.total} : < :"""
										elif self.predicted_remaining is None:
																snake_case       :     List[str] =    F"""[{spaced_value}/{self.total} {format_time(self.elapsed_time      )}"""
										else:
																snake_case       :     Optional[int] =    (
																    F"""[{spaced_value}/{self.total} {format_time(self.elapsed_time      )} <"""
																    F""" {format_time(self.predicted_remaining      )}"""
																)
																self.label += F""", {1/self.average_time_per_item:.2f} it/s"""
										self.label += "]" if self.comment is None or len(self.comment      ) == 0 else F""", {self.comment}]"""
										self.display()




				def 	__snake_case							(							self	:				Optional[int]      ):
										snake_case       :     str =    html_progress_bar(self.value,    self.total,    self.prefix,    self.label,    self.width      )
										if self.parent is not None:
																# If this is a child bar, the parent will take care of the display.
																self.parent.display()
																return
										if self.output is None:
																snake_case       :     Any =    disp.display(disp.HTML(self.html_code      ),    display_id=SCREAMING_SNAKE_CASE_      )
										else:
																self.output.update(disp.HTML(self.html_code      )      )




				def 	__snake_case							(							self	:				Optional[int]      ):
										if self.parent is None and self.output is not None:
																self.output.update(disp.HTML(''''''      )      )




class 						a   (				__magic_name__    ):




				def __init__(							self	:				List[str],    SCREAMING_SNAKE_CASE_	:				List[str],    SCREAMING_SNAKE_CASE_	:				Optional[int]=None      ):
										super().__init__(SCREAMING_SNAKE_CASE_      )
										snake_case       :     Tuple =    None if column_names is None else [column_names]
										snake_case       :     str =    None




				def 	__snake_case							(							self	:				Dict      ):
										snake_case       :     Tuple =    html_progress_bar(self.value,    self.total,    self.prefix,    self.label,    self.width      )
										if self.inner_table is not None:
																self.html_code += text_to_html_table(self.inner_table      )
										if self.child_bar is not None:
																self.html_code += self.child_bar.html_code
										if self.output is None:
																snake_case       :     List[Any] =    disp.display(disp.HTML(self.html_code      ),    display_id=SCREAMING_SNAKE_CASE_      )
										else:
																self.output.update(disp.HTML(self.html_code      )      )




				def 	__snake_case							(							self	:				List[str],    SCREAMING_SNAKE_CASE_	:				Any      ):
										if self.inner_table is None:
																snake_case       :     Optional[Any] =    [list(values.keys()      ), list(values.values()      )]
										else:
																snake_case       :     Tuple =    self.inner_table[0]
																if len(self.inner_table      ) == 1:
																						# We give a chance to update the column names at the first iteration
																						for key in values.keys():
																												if key not in columns:
																																		columns.append(SCREAMING_SNAKE_CASE_      )
																						snake_case       :     int =    columns
																self.inner_table.append([values[c] for c in columns]      )




				def 	__snake_case							(							self	:				Tuple,    SCREAMING_SNAKE_CASE_	:				Optional[Any],    SCREAMING_SNAKE_CASE_	:				int=None,    SCREAMING_SNAKE_CASE_	:				str=3_00      ):
										snake_case       :     int =    NotebookProgressBar(SCREAMING_SNAKE_CASE_,    prefix=SCREAMING_SNAKE_CASE_,    parent=self,    width=SCREAMING_SNAKE_CASE_      )
										return self.child_bar




				def 	__snake_case							(							self	:				Union[str, Any]      ):
										snake_case       :     int =    None
										self.display()




class 						a   (				__magic_name__    ):




				def __init__(							self	:				Optional[Any]      ):
										snake_case       :     Dict =    None
										snake_case       :     Union[str, Any] =    None
										snake_case       :     Dict =    False




				def 	__snake_case							(							self	:				List[str],    SCREAMING_SNAKE_CASE_	:				Any,    SCREAMING_SNAKE_CASE_	:				int,    SCREAMING_SNAKE_CASE_	:				str,    **SCREAMING_SNAKE_CASE_	:				List[str]      ):
										snake_case       :     str =    '''Epoch''' if args.evaluation_strategy == IntervalStrategy.EPOCH else '''Step'''
										snake_case       :     List[str] =    0
										snake_case       :     Dict =    0
										snake_case       :     List[Any] =    [self.first_column] + ['''Training Loss''']
										if args.evaluation_strategy != IntervalStrategy.NO:
																column_names.append('''Validation Loss'''      )
										snake_case       :     Tuple =    NotebookTrainingTracker(state.max_steps,    SCREAMING_SNAKE_CASE_      )




				def 	__snake_case							(							self	:				Tuple,    SCREAMING_SNAKE_CASE_	:				List[Any],    SCREAMING_SNAKE_CASE_	:				Dict,    SCREAMING_SNAKE_CASE_	:				str,    **SCREAMING_SNAKE_CASE_	:				Union[str, Any]      ):
										snake_case       :     Tuple =    int(state.epoch      ) if int(state.epoch      ) == state.epoch else F"""{state.epoch:.2f}"""
										self.training_tracker.update(
										    state.global_step + 1,    comment=F"""Epoch {epoch}/{state.num_train_epochs}""",    force_update=self._force_next_update,    )
										snake_case       :     str =    False




				def 	__snake_case							(							self	:				Tuple,    SCREAMING_SNAKE_CASE_	:				int,    SCREAMING_SNAKE_CASE_	:				Union[str, Any],    SCREAMING_SNAKE_CASE_	:				int,    SCREAMING_SNAKE_CASE_	:				Tuple=None,    **SCREAMING_SNAKE_CASE_	:				Tuple      ):
										if not has_length(SCREAMING_SNAKE_CASE_      ):
																return
										if self.prediction_bar is None:
																if self.training_tracker is not None:
																						snake_case       :     Dict =    self.training_tracker.add_child(len(SCREAMING_SNAKE_CASE_      )      )
																else:
																						snake_case       :     Optional[Any] =    NotebookProgressBar(len(SCREAMING_SNAKE_CASE_      )      )
																self.prediction_bar.update(1      )
										else:
																self.prediction_bar.update(self.prediction_bar.value + 1      )




				def 	__snake_case							(							self	:				Optional[int],    SCREAMING_SNAKE_CASE_	:				Any,    SCREAMING_SNAKE_CASE_	:				Dict,    SCREAMING_SNAKE_CASE_	:				Union[str, Any],    **SCREAMING_SNAKE_CASE_	:				List[Any]      ):
										if self.prediction_bar is not None:
																self.prediction_bar.close()
										snake_case       :     Optional[int] =    None




				def 	__snake_case							(							self	:				Optional[Any],    SCREAMING_SNAKE_CASE_	:				List[str],    SCREAMING_SNAKE_CASE_	:				Any,    SCREAMING_SNAKE_CASE_	:				Optional[int],    SCREAMING_SNAKE_CASE_	:				Optional[int]=None,    **SCREAMING_SNAKE_CASE_	:				Union[str, Any]      ):
										# Only for when there is no evaluation
										if args.evaluation_strategy == IntervalStrategy.NO and "loss" in logs:
																snake_case       :     List[str] =    {'''Training Loss''': logs['''loss''']}
																# First column is necessarily Step sine we're not in epoch eval strategy
																snake_case       :     List[Any] =    state.global_step
																self.training_tracker.write_line(SCREAMING_SNAKE_CASE_      )




				def 	__snake_case							(							self	:				Any,    SCREAMING_SNAKE_CASE_	:				Optional[int],    SCREAMING_SNAKE_CASE_	:				List[str],    SCREAMING_SNAKE_CASE_	:				List[str],    SCREAMING_SNAKE_CASE_	:				Optional[Any]=None,    **SCREAMING_SNAKE_CASE_	:				Any      ):
										if self.training_tracker is not None:
																snake_case       :     List[Any] =    {'''Training Loss''': '''No log''', '''Validation Loss''': '''No log'''}
																for log in reversed(state.log_history      ):
																						if "loss" in log:
																												snake_case       :     Union[str, Any] =    log['''loss''']
																												break

																if self.first_column == "Epoch":
																						snake_case       :     Optional[int] =    int(state.epoch      )
																else:
																						snake_case       :     Optional[int] =    state.global_step
																snake_case       :     Optional[int] =    '''eval'''
																for k in metrics:
																						if k.endswith('''_loss'''      ):
																												snake_case       :     str =    re.sub(R'''\_loss$''',    '''''',    SCREAMING_SNAKE_CASE_      )
																snake_case       :     str =    metrics.pop('''total_flos''',    SCREAMING_SNAKE_CASE_      )
																snake_case       :     int =    metrics.pop('''epoch''',    SCREAMING_SNAKE_CASE_      )
																snake_case       :     List[str] =    metrics.pop(F"""{metric_key_prefix}_runtime""",    SCREAMING_SNAKE_CASE_      )
																snake_case       :     Union[str, Any] =    metrics.pop(F"""{metric_key_prefix}_samples_per_second""",    SCREAMING_SNAKE_CASE_      )
																snake_case       :     Any =    metrics.pop(F"""{metric_key_prefix}_steps_per_second""",    SCREAMING_SNAKE_CASE_      )
																snake_case       :     Dict =    metrics.pop(F"""{metric_key_prefix}_jit_compilation_time""",    SCREAMING_SNAKE_CASE_      )
																for k, v in metrics.items():
																						if k == F"""{metric_key_prefix}_loss""":
																												snake_case       :     Optional[int] =    v
																						else:
																												snake_case       :     int =    k.split('''_'''      )
																												snake_case       :     Any =    ''' '''.join([part.capitalize() for part in splits[1:]]      )
																												snake_case       :     Tuple =    v
																self.training_tracker.write_line(SCREAMING_SNAKE_CASE_      )
																self.training_tracker.remove_child()
																snake_case       :     Any =    None
																# Evaluation takes a long time so we should force the next update.
																snake_case       :     Optional[int] =    True




				def 	__snake_case							(							self	:				List[str],    SCREAMING_SNAKE_CASE_	:				List[Any],    SCREAMING_SNAKE_CASE_	:				List[str],    SCREAMING_SNAKE_CASE_	:				Dict,    **SCREAMING_SNAKE_CASE_	:				str      ):
										self.training_tracker.update(
										    state.global_step,    comment=F"""Epoch {int(state.epoch      )}/{state.num_train_epochs}""",    force_update=SCREAMING_SNAKE_CASE_      )
										snake_case       :     Tuple =    None

