import argparse
import shutil
import time
from json import JSONDecodeError
from logging import getLogger
from pathlib import Path
from typing import Dict, List

import torch
from torch.utils.data import DataLoader
from tqdm import tqdm

from transformers import AutoModelForSeqaSeqLM, AutoTokenizer
from utils import (
    SeqaSeqDataset,
    calculate_bleu,
    calculate_rouge,
    chunks,
    lmap,
    load_json,
    parse_numeric_n_bool_cl_kwargs,
    save_json,
    use_task_specific_params,
    write_txt_file,
)


_SCREAMING_SNAKE_CASE     				=		getLogger(__name__)






def 					snake_case    (			snake_case__     :Optional[Any]  ,      snake_case__     :Dict  ,      snake_case__     :List[Any]  ,      snake_case__     :List[Any] = 8  ,      snake_case__     :Optional[int] = 1_024  ,      snake_case__     :Dict="val"  ,      snake_case__     :Dict=None  ,      snake_case__     :int=False  ,      snake_case__     :str="summarization"  ,      snake_case__     :str=None  ,      snake_case__     :Union[str, Any]=1  ,      snake_case__     :str = None  ,      snake_case__     :int=""  ,      **snake_case__     :Any  ,      )		->  Dict:
       _A	  =    str(__a)
       assert local_rank is not None
       torch.distributed.init_process_group(backend="""nccl"""  ,      rank=__a)

       _A	  =    Path(__a)
       _A	  =    save_dir.joinpath(F'''rank_{local_rank}_output.json''')
       torch.cuda.set_device(__a)
       _A	  =    AutoModelForSeqaSeqLM.from_pretrained(__a).cuda()
       if fpaa:
              _A	  =    model.half()
       # determine if we need to increase num_beams
       use_task_specific_params(__a  ,      __a)  # update config with task specific params
       _A	  =    generate_kwargs.pop("""num_beams"""  ,      model.config.num_beams)  # AttributeError risk?
       if num_return_sequences > num_beams:
              _A	  =    num_return_sequences

       _A	  =    AutoTokenizer.from_pretrained(__a)
       logger.info(F'''Inferred tokenizer type: {tokenizer.__class__}''')  # if this is wrong, check config.model_type.

       if max_source_length is None:
              _A	  =    tokenizer.model_max_length
       if prefix is None:
              _A	  =    prefix or getattr(model.config  ,      """prefix"""  ,      """""") or ''''''
       _A	  =    SeqaSeqDataset(
           __a  ,      __a  ,      __a  ,      max_target_length=1_024  ,      type_path=__a  ,      n_obs=__a  ,      prefix=__a  ,      **__a  ,      )
       # I set shuffle=True for a more accurate progress bar.
       # If all the longest samples are first, the prog bar estimate is too high at the beginning.
       _A	  =    ds.make_sortish_sampler(__a  ,      distributed=__a  ,      add_extra_examples=__a  ,      shuffle=__a)
       _A	  =    DataLoader(__a  ,      sampler=__a  ,      batch_size=__a  ,      collate_fn=ds.collate_fn)
       _A	  =    []
       for batch in tqdm(__a):
              _A	  =    model.generate(
                  input_ids=batch["""input_ids"""].to(model.device)  ,      attention_mask=batch["""attention_mask"""].to(model.device)  ,      num_return_sequences=__a  ,      num_beams=__a  ,      **__a  ,      )
              _A	  =    tokenizer.batch_decode(__a  ,      skip_special_tokens=__a  ,      clean_up_tokenization_spaces=__a)
              _A	  =    batch['''ids''']
              if num_return_sequences > 1:
                     _A	  =    chunks(__a  ,      __a)  # batch size chunks, each of size num_return_seq
              for i, pred in enumerate(__a):
                     results.append({"""pred""": pred, """id""": ids[i].item()})
       save_json(__a  ,      __a)
       return results, sampler.num_replicas






def 					snake_case    (			)		->  int:
       _A	  =    argparse.ArgumentParser(
           epilog="""Unspecified args like --num_beams=2 --decoder_start_token_id=4 are passed to model.generate""")
       parser.add_argument("""--data_dir"""  ,      type=__a  ,      help="""like cnn_dm/test.source""")
       parser.add_argument(
           """--model_name"""  ,      type=__a  ,      help="""like facebook/bart-large-cnn,t5-base, etc."""  ,      default="""sshleifer/distilbart-xsum-12-3"""  ,      )
       parser.add_argument("""--save_dir"""  ,      type=__a  ,      help="""where to save"""  ,      default="""tmp_gen""")
       parser.add_argument("""--max_source_length"""  ,      type=__a  ,      default=__a)
       parser.add_argument(
           """--type_path"""  ,      type=__a  ,      default="""test"""  ,      help="""which subset to evaluate typically train/val/test""")
       parser.add_argument("""--task"""  ,      type=__a  ,      default="""summarization"""  ,      help="""used for task_specific_params + metrics""")
       parser.add_argument("""--bs"""  ,      type=__a  ,      default=8  ,      required=__a  ,      help="""batch size""")
       parser.add_argument(
           """--local_rank"""  ,      type=__a  ,      default=-1  ,      required=__a  ,      help="""should be passed by distributed.launch""")

       parser.add_argument(
           """--n_obs"""  ,      type=__a  ,      default=__a  ,      required=__a  ,      help="""How many observations. Defaults to all.""")
       parser.add_argument(
           """--num_return_sequences"""  ,      type=__a  ,      default=1  ,      required=__a  ,      help="""How many sequences to return""")
       parser.add_argument(
           """--sync_timeout"""  ,      type=__a  ,      default=600  ,      required=__a  ,      help="""How long should master process wait for other processes to finish."""  ,      )
       parser.add_argument("""--src_lang"""  ,      type=__a  ,      default=__a  ,      required=__a)
       parser.add_argument("""--tgt_lang"""  ,      type=__a  ,      default=__a  ,      required=__a)
       parser.add_argument(
           """--prefix"""  ,      type=__a  ,      required=__a  ,      default=__a  ,      help="""will be added to the begininng of src examples""")
       parser.add_argument("""--fp16"""  ,      action="""store_true""")
       parser.add_argument("""--debug"""  ,      action="""store_true""")
       _A	  =    time.time()
       _A	  =    parser.parse_known_args()
       _A	  =    parse_numeric_n_bool_cl_kwargs(__a)
       if generate_kwargs and args.local_rank <= 0:
              print(F'''parsed the following generate kwargs: {generate_kwargs}''')
       _A	  =    Path(args.save_dir + """_tmp""")
       Path(__a).mkdir(exist_ok=__a)  # this handles locking.
       _A	  =    list(json_save_dir.glob("""rank_*.json"""))
       if intermediate_files:
              raise ValueError(F'''Found files at {json_save_dir} please move or remove them.''')
              # In theory, a node could finish and save before another node hits this. If this happens, we can address later.
       _A	  =    {}
       if args.src_lang is not None:
              _A	  =    args.src_lang
       if args.tgt_lang is not None:
              _A	  =    args.tgt_lang

       Path(args.save_dir).mkdir(exist_ok=__a)
       _A	  =    eval_data_dir(
           args.data_dir  ,      __a  ,      args.model_name  ,      type_path=args.type_path  ,      bs=args.bs  ,      fpaa=args.fpaa  ,      task=args.task  ,      local_rank=args.local_rank  ,      n_obs=args.n_obs  ,      max_source_length=args.max_source_length  ,      num_return_sequences=args.num_return_sequences  ,      prefix=args.prefix  ,      dataset_kwargs=__a  ,      **__a  ,      )

       if args.local_rank <= 0:
              _A	  =    Path(args.save_dir)
              save_dir.mkdir(exist_ok=__a)
              _A	  =    gather_results_from_each_node(__a  ,      __a  ,      args.sync_timeout)
              _A	  =    combine_partial_results(__a)
              if args.num_return_sequences > 1:
                     _A	  =    save_dir.joinpath("""pseudolabel_results.json""")
                     print(F'''Saving aggregated results at {save_path}, intermediate in {json_save_dir}/''')
                     save_json(__a  ,      __a)
                     return
              _A	  =    Path(args.data_dir).joinpath(args.type_path + """.target""")
              with open(__a) as f:
                     _A	  =    [x.rstrip() for x in f.readlines()][: len(__a)]

              # Calculate metrics, save metrics,  and save _generations.txt
              _A	  =    '''translation''' in args.task
              _A	  =    calculate_bleu if calc_bleu else calculate_rouge
              _A	  =    '''bleu''' if calc_bleu else '''rouge'''
              _A	  =    score_fn(__a  ,      __a)
              _A	  =    len(__a)
              _A	  =    time.time() - start_time
              _A	  =    round(runtime / metrics["""n_obs"""]  ,      4)
              _A	  =    num_replicas
              # TODO(@stas00): add whatever metadata to metrics
              _A	  =    save_dir.joinpath(F'''{args.type_path}_{metric_name}.json''')
              save_json(__a  ,      __a  ,      indent=__a)
              print(__a)
              write_txt_file(__a  ,      save_dir.joinpath(F'''{args.type_path}_generations.txt'''))
              if args.debug:
                     write_txt_file(__a  ,      save_dir.joinpath(F'''{args.type_path}.target'''))
              else:
                     shutil.rmtree(__a)






def 					snake_case    (			snake_case__     :Dict)		->  List:
       _A	  =    []
       for partial_result in partial_results:
              records.extend(__a)
       _A	  =    sorted(__a  ,      key=lambda snake_case__: x["id"])
       _A	  =    [x['''pred'''] for x in records]
       return preds






def 					snake_case    (			snake_case__     :List[Any]  ,      snake_case__     :Tuple  ,      snake_case__     :Tuple)		->  List[Dict[str, List]]:
       _A	  =    time.time()
       logger.info("""waiting for all nodes to finish""")
       _A	  =    None
       while (time.time() - start_wait) < timeout:
              _A	  =    list(save_dir.glob("""rank_*.json"""))
              if len(__a) < num_replicas:
                     continue
              try:
                     # make sure all json files are fully saved
                     _A	  =    lmap(__a  ,      __a)
                     return json_data
              except JSONDecodeError:
                     continue
       else:
              raise TimeoutError("""Rank 0 gave up on waiting for other processes""")
       # Unreachable


if __name__ == "__main__":
    # Usage for MT:
    run_generate()








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

		"""simple docstring"""

		def __init__(						self       ,  UpperCamelCase_       ,  UpperCamelCase_=3       ,  UpperCamelCase_=7       ,  UpperCamelCase_=True       ,  UpperCamelCase_=True       ,  UpperCamelCase_=False       ,  UpperCamelCase_=True       ,  UpperCamelCase_=99       ,  UpperCamelCase_=32       ,  UpperCamelCase_=5       ,  UpperCamelCase_=4       ,  UpperCamelCase_=37       ,  UpperCamelCase_="gelu"       ,  UpperCamelCase_=0.1       ,  UpperCamelCase_=0.1       ,  UpperCamelCase_=512       ,  UpperCamelCase_=16       ,  UpperCamelCase_=2       ,  UpperCamelCase_=0.02       ,  UpperCamelCase_=3       ,  UpperCamelCase_=4       ,  UpperCamelCase_=None       ,  ):

							'''simple docstring'''


							UpperCamelCase__  :Dict				   =					parent
							UpperCamelCase__  :Any				   =					batch_size
							UpperCamelCase__  :Tuple				   =					seq_length
							UpperCamelCase__  :str				   =					is_training
							UpperCamelCase__  :str				   =					use_input_mask
							UpperCamelCase__  :Union[str, Any]				   =					use_token_type_ids
							UpperCamelCase__  :int				   =					use_labels
							UpperCamelCase__  :Union[str, Any]				   =					vocab_size
							UpperCamelCase__  :int				   =					hidden_size
							UpperCamelCase__  :Union[str, Any]				   =					num_hidden_layers
							UpperCamelCase__  :List[Any]				   =					num_attention_heads
							UpperCamelCase__  :int				   =					intermediate_size
							UpperCamelCase__  :str				   =					hidden_act
							UpperCamelCase__  :List[str]				   =					hidden_dropout_prob
							UpperCamelCase__  :Tuple				   =					attention_probs_dropout_prob
							UpperCamelCase__  :str				   =					max_position_embeddings
							UpperCamelCase__  :List[Any]				   =					type_vocab_size
							UpperCamelCase__  :Any				   =					type_sequence_label_size
							UpperCamelCase__  :Dict				   =					initializer_range
							UpperCamelCase__  :List[str]				   =					num_labels
							UpperCamelCase__  :Tuple				   =					num_choices
							UpperCamelCase__  :List[Any]				   =					scope




		def        lowerCAmelCase__							(						self				):

							'''simple docstring'''


							UpperCamelCase__  :Any				   =					ids_tensor([self.batch_size, self.seq_length]       ,  self.vocab_size				)

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

							UpperCamelCase__  :Optional[Any]				   =					None

							UpperCamelCase__  :Optional[Any]				   =					None
							UpperCamelCase__  :Optional[int]				   =					None
							UpperCamelCase__  :Optional[int]				   =					None
							if self.use_labels:
												UpperCamelCase__  :Optional[Any]				   =					ids_tensor([self.batch_size]       ,  self.type_sequence_label_size				)
												UpperCamelCase__  :Any				   =					ids_tensor([self.batch_size, self.seq_length]       ,  self.num_labels				)
												UpperCamelCase__  :Any				   =					ids_tensor([self.batch_size]       ,  self.num_choices				)

							UpperCamelCase__  :Any				   =					self.get_config()

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




		def        lowerCAmelCase__							(						self				):

							'''simple docstring'''


							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       ,  UpperCamelCase_       ,  UpperCamelCase_       ,  UpperCamelCase_       ,  UpperCamelCase_       ,  UpperCamelCase_       ,  UpperCamelCase_       ,  UpperCamelCase_				):

							'''simple docstring'''


							UpperCamelCase__  :Any				   =					FalconModel(config=UpperCamelCase_				)
							model.to(UpperCamelCase_				)
							model.eval()
							UpperCamelCase__  :Dict				   =					model(UpperCamelCase_       ,  attention_mask=UpperCamelCase_				)
							UpperCamelCase__  :Tuple				   =					model(UpperCamelCase_				)
							self.parent.assertEqual(result.last_hidden_state.shape       ,  (self.batch_size, self.seq_length, self.hidden_size)				)




		def        lowerCAmelCase__							(						self       ,  UpperCamelCase_       ,  UpperCamelCase_       ,  UpperCamelCase_       ,  UpperCamelCase_       ,  UpperCamelCase_       ,  UpperCamelCase_       ,  UpperCamelCase_       ,  UpperCamelCase_       ,  UpperCamelCase_       ,  ):

							'''simple docstring'''


							UpperCamelCase__  :Union[str, Any]				   =					True
							UpperCamelCase__  :Dict				   =					FalconModel(UpperCamelCase_				)
							model.to(UpperCamelCase_				)
							model.eval()
							UpperCamelCase__  :Optional[Any]				   =					model(
							    UpperCamelCase_       ,  attention_mask=UpperCamelCase_       ,  encoder_hidden_states=UpperCamelCase_       ,  encoder_attention_mask=UpperCamelCase_       ,  )
							UpperCamelCase__  :List[str]				   =					model(
							    UpperCamelCase_       ,  attention_mask=UpperCamelCase_       ,  encoder_hidden_states=UpperCamelCase_       ,  )
							UpperCamelCase__  :str				   =					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       ,  UpperCamelCase_       ,  UpperCamelCase_       ,  UpperCamelCase_       ,  UpperCamelCase_       ,  UpperCamelCase_       ,  UpperCamelCase_       ,  UpperCamelCase_       ,  UpperCamelCase_       ,  UpperCamelCase_       ,  ):

							'''simple docstring'''


							UpperCamelCase__  :Tuple				   =					FalconForCausalLM(config=UpperCamelCase_				)
							model.to(UpperCamelCase_				)
							model.eval()
							UpperCamelCase__  :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       ,  UpperCamelCase_       ,  UpperCamelCase_       ,  UpperCamelCase_       ,  UpperCamelCase_       ,  UpperCamelCase_       ,  UpperCamelCase_       ,  UpperCamelCase_       ,  UpperCamelCase_       ,  UpperCamelCase_       ,  ):

							'''simple docstring'''


							UpperCamelCase__  :str				   =					True
							UpperCamelCase__  :Tuple				   =					True
							UpperCamelCase__  :Any				   =					FalconForCausalLM(config=UpperCamelCase_				)
							model.to(UpperCamelCase_				)
							model.eval()

							# first forward pass
							UpperCamelCase__  :Dict				   =					model(
							    UpperCamelCase_       ,  attention_mask=UpperCamelCase_       ,  encoder_hidden_states=UpperCamelCase_       ,  encoder_attention_mask=UpperCamelCase_       ,  use_cache=UpperCamelCase_       ,  )
							UpperCamelCase__  :List[Any]				   =					outputs.past_key_values

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

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

							UpperCamelCase__  :Any				   =					model(
							    UpperCamelCase_       ,  attention_mask=UpperCamelCase_       ,  encoder_hidden_states=UpperCamelCase_       ,  encoder_attention_mask=UpperCamelCase_       ,  output_hidden_states=UpperCamelCase_       ,  )['''hidden_states'''][0]
							UpperCamelCase__  :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
							UpperCamelCase__  :List[Any]				   =					ids_tensor((1,)       ,  output_from_past.shape[-1]				).item()
							UpperCamelCase__  :Optional[int]				   =					output_from_no_past[:, -3:, random_slice_idx].detach()
							UpperCamelCase__  :Tuple				   =					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				):

							'''simple docstring'''


							UpperCamelCase__  :Dict				   =					self.prepare_config_and_inputs()
							(
							    (
							    UpperCamelCase__
							)   ,							(
							    UpperCamelCase__
							)   ,							(
							    UpperCamelCase__
							)   ,							(
							    UpperCamelCase__
							)   ,							(
							    UpperCamelCase__
							)   ,							(
							    UpperCamelCase__
							)   ,							(
							    UpperCamelCase__
							)   ,							
							)  :Optional[int]				   =					config_and_inputs
							UpperCamelCase__  :Optional[int]				   =					{'''input_ids''': input_ids, '''attention_mask''': input_mask}
							return config, inputs_dict



@require_torch
class   lowercase      (				A__      ,						A__      ,						A__      ,						unittest.TestCase   ):

		"""simple docstring"""

		_a       					=							(
		    (
		        FalconModel,
		        FalconForCausalLM,
		        FalconForSequenceClassification,
		        FalconForTokenClassification,
		        FalconForQuestionAnswering,
		    )
		    if is_torch_available()
		    else ()
		)
		_a       					=							(FalconForCausalLM,) if is_torch_available() else ()
		_a       					=							(
		    {
		        'feature-extraction': FalconModel,
		        'text-classification': FalconForSequenceClassification,
		        'text-generation': FalconForCausalLM,
		        'question-answering': FalconForQuestionAnswering,
		        'token-classification': FalconForTokenClassification,
		        'zero-shot': FalconForSequenceClassification,
		    }
		    if is_torch_available()
		    else {}
		)
		_a       					=							False
		_a       					=							False




		def        lowerCAmelCase__							(						self				):

							'''simple docstring'''


							UpperCamelCase__  :Union[str, Any]				   =					FalconModelTester(self				)
							UpperCamelCase__  :List[Any]				   =					ConfigTester(self       ,  config_class=UpperCamelCase_       ,  hidden_size=37				)




		def        lowerCAmelCase__							(						self				):

							'''simple docstring'''


							self.config_tester.run_common_tests()




		def        lowerCAmelCase__							(						self				):

							'''simple docstring'''


							UpperCamelCase__  :Tuple				   =					self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*UpperCamelCase_				)




		def        lowerCAmelCase__							(						self				):

							'''simple docstring'''


							UpperCamelCase__   ,							*UpperCamelCase__  :Union[str, Any]				   =					self.model_tester.prepare_config_and_inputs()
							for alibi in [True, False]:
												UpperCamelCase__  :Any				   =					alibi
												self.model_tester.create_and_check_model(UpperCamelCase_       ,  *UpperCamelCase_				)




		def        lowerCAmelCase__							(						self				):

							'''simple docstring'''


							UpperCamelCase__   ,							UpperCamelCase__  :List[str]				   =					self.model_tester.prepare_config_and_inputs_for_common()
							UpperCamelCase__  :str				   =					3
							UpperCamelCase__  :Any				   =					input_dict['''input_ids''']
							UpperCamelCase__  :Tuple				   =					input_ids.ne(1				).to(UpperCamelCase_				)
							UpperCamelCase__  :Optional[Any]				   =					ids_tensor([self.model_tester.batch_size]       ,  self.model_tester.type_sequence_label_size				)
							UpperCamelCase__  :List[Any]				   =					FalconForSequenceClassification(UpperCamelCase_				)
							model.to(UpperCamelCase_				)
							model.eval()
							UpperCamelCase__  :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				):

							'''simple docstring'''


							UpperCamelCase__   ,							UpperCamelCase__  :List[Any]				   =					self.model_tester.prepare_config_and_inputs_for_common()
							UpperCamelCase__  :List[str]				   =					3
							UpperCamelCase__  :str				   =					'''single_label_classification'''
							UpperCamelCase__  :Union[str, Any]				   =					input_dict['''input_ids''']
							UpperCamelCase__  :List[str]				   =					input_ids.ne(1				).to(UpperCamelCase_				)
							UpperCamelCase__  :Dict				   =					ids_tensor([self.model_tester.batch_size]       ,  self.model_tester.type_sequence_label_size				)
							UpperCamelCase__  :Dict				   =					FalconForSequenceClassification(UpperCamelCase_				)
							model.to(UpperCamelCase_				)
							model.eval()
							UpperCamelCase__  :str				   =					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				):

							'''simple docstring'''


							UpperCamelCase__   ,							UpperCamelCase__  :Tuple				   =					self.model_tester.prepare_config_and_inputs_for_common()
							UpperCamelCase__  :int				   =					input_dict['''input_ids''']
							UpperCamelCase__  :List[str]				   =					FalconForCausalLM(UpperCamelCase_				)
							model.to(UpperCamelCase_				)
							model.eval()
							UpperCamelCase__  :Optional[int]				   =					model(UpperCamelCase_       ,  use_cache=UpperCamelCase_				)
							UpperCamelCase__  :Any				   =					input_ids.shape[0]
							UpperCamelCase__  :Optional[int]				   =					model._convert_to_rw_cache(result.past_key_values				)
							UpperCamelCase__  :List[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				):

							'''simple docstring'''


							UpperCamelCase__   ,							UpperCamelCase__  :List[Any]				   =					self.model_tester.prepare_config_and_inputs_for_common()
							UpperCamelCase__  :Optional[int]				   =					3
							UpperCamelCase__  :Union[str, Any]				   =					'''multi_label_classification'''
							UpperCamelCase__  :Dict				   =					input_dict['''input_ids''']
							UpperCamelCase__  :Tuple				   =					input_ids.ne(1				).to(UpperCamelCase_				)
							UpperCamelCase__  :Union[str, Any]				   =					ids_tensor(
							    [self.model_tester.batch_size, config.num_labels]       ,  self.model_tester.type_sequence_label_size				).to(torch.float				)
							UpperCamelCase__  :Optional[Any]				   =					FalconForSequenceClassification(UpperCamelCase_				)
							model.to(UpperCamelCase_				)
							model.eval()
							UpperCamelCase__  :str				   =					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				):

							'''simple docstring'''


							for model_class in self.all_generative_model_classes:
												UpperCamelCase__   ,							UpperCamelCase__  :List[Any]				   =					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

												UpperCamelCase__  :Tuple				   =					model_class(UpperCamelCase_				).to(UpperCamelCase_				)
												if "use_cache" not in inputs:
																	UpperCamelCase__  :Tuple				   =					True
												UpperCamelCase__  :List[str]				   =					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

												UpperCamelCase__  :str				   =					(
												    getattr(UpperCamelCase_       ,  '''decoder_layers'''       ,  UpperCamelCase_				)
												    or getattr(UpperCamelCase_       ,  '''num_decoder_layers'''       ,  UpperCamelCase_				)
												    or config.num_hidden_layers
												)
												UpperCamelCase__  :List[Any]				   =					getattr(UpperCamelCase_       ,  '''num_kv_heads'''       ,  config.num_attention_heads				)
												UpperCamelCase__  :str				   =					getattr(UpperCamelCase_       ,  '''d_model'''       ,  config.hidden_size				)
												UpperCamelCase__  :Optional[Any]				   =					embed_dim // num_attention_heads

												UpperCamelCase__  :Optional[Any]				   =					outputs['''past_key_values''']
												self.assertEqual(len(UpperCamelCase_				)       ,  UpperCamelCase_				)

												UpperCamelCase__   ,							UpperCamelCase__  :Tuple				   =					inputs['''input_ids'''].shape
												for i in range(UpperCamelCase_				):
																	if config.new_decoder_architecture:
																						UpperCamelCase__  :List[Any]				   =					config.num_attention_heads
																	elif config.multi_query:
																						UpperCamelCase__  :str				   =					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   lowercase      (				unittest.TestCase   ):

		"""simple docstring"""

		@slow
		def        lowerCAmelCase__							(						self				):

							'''simple docstring'''


							UpperCamelCase__  :List[str]				   =					AutoTokenizer.from_pretrained('''Rocketknight1/falcon-rw-1b'''				)
							UpperCamelCase__  :List[Any]				   =					FalconForCausalLM.from_pretrained('''Rocketknight1/falcon-rw-1b'''				)
							model.eval()
							model.to(UpperCamelCase_				)
							UpperCamelCase__  :List[str]				   =					tokenizer('''My favorite food is'''       ,  return_tensors='''pt'''				).to(UpperCamelCase_				)

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

							UpperCamelCase__  :int				   =					model.generate(**UpperCamelCase_       ,  do_sample=UpperCamelCase_       ,  max_new_tokens=19				)
							UpperCamelCase__  :str				   =					tokenizer.batch_decode(UpperCamelCase_				)[0]

							self.assertEqual(UpperCamelCase_       ,  UpperCamelCase_				)




		@slow
		def        lowerCAmelCase__							(						self				):

							'''simple docstring'''


							for repo in ["Rocketknight1/tiny-random-falcon-7b", "Rocketknight1/tiny-random-falcon-40b"]:
												UpperCamelCase__  :Any				   =					AutoTokenizer.from_pretrained(UpperCamelCase_				)
												UpperCamelCase__  :Dict				   =					FalconForCausalLM.from_pretrained(UpperCamelCase_				)
												model.eval()
												model.to(UpperCamelCase_				)
												UpperCamelCase__  :str				   =					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				):

							'''simple docstring'''


							with torch.no_grad():
												for repo in [
												    "Rocketknight1/falcon-rw-1b",
												    "Rocketknight1/tiny-random-falcon-7b",
												    "Rocketknight1/tiny-random-falcon-40b",
												]:
																	UpperCamelCase__  :List[Any]				   =					AutoTokenizer.from_pretrained(UpperCamelCase_				)
																	UpperCamelCase__  :Union[str, Any]				   =					FalconForCausalLM.from_pretrained(UpperCamelCase_				)
																	model.eval()
																	model.to(device=UpperCamelCase_				)
																	UpperCamelCase__  :Tuple				   =					tokenizer('''My favorite food is'''       ,  return_tensors='''pt'''				).to(UpperCamelCase_				)

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