



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    :				Optional[int]       ,	snake_case    :				List[str]       ,	snake_case    :				Union[str, Any]=3       ,	snake_case    :				List[Any]=7       ,	snake_case    :				List[str]=True       ,	snake_case    :				Tuple=True       ,	snake_case    :				Any=False       ,	snake_case    :				List[str]=True       ,	snake_case    :				List[str]=99       ,	snake_case    :				Tuple=32       ,	snake_case    :				Union[str, Any]=5       ,	snake_case    :				str=4       ,	snake_case    :				Dict=37       ,	snake_case    :				str="gelu"       ,	snake_case    :				List[Any]=0.1       ,	snake_case    :				str=0.1       ,	snake_case    :				Union[str, Any]=512       ,	snake_case    :				Optional[int]=16       ,	snake_case    :				Tuple=2       ,	snake_case    :				List[str]=0.02       ,	snake_case    :				int=3       ,	snake_case    :				Optional[Any]=4       ,	snake_case    :				int=None       ,	):



											'''simple docstring'''


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







						def       lowerCamelCase_				(  self    :				Optional[Any]	):



											'''simple docstring'''


											SCREAMING_SNAKE_CASE							:     str						=				ids_tensor([self.batch_size, self.seq_length]       ,	self.vocab_size	)

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

											SCREAMING_SNAKE_CASE							:     List[Any]						=				None

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

											SCREAMING_SNAKE_CASE							:     Dict						=				self.get_config()

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







						def       lowerCamelCase_				(  self    :				List[Any]	):



											'''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=snake_case       ,	initializer_range=self.initializer_range       ,	pad_token_id=1       ,	new_decoder_architecture=snake_case       ,	)







						def       lowerCamelCase_				(  self    :				List[str]       ,	snake_case    :				Optional[Any]       ,	snake_case    :				Union[str, Any]       ,	snake_case    :				Any       ,	snake_case    :				Any       ,	snake_case    :				int       ,	snake_case    :				List[str]       ,	snake_case    :				Tuple	):



											'''simple docstring'''


											SCREAMING_SNAKE_CASE							:     str						=				FalconModel(config=snake_case	)
											model.to(snake_case	)
											model.eval()
											SCREAMING_SNAKE_CASE							:     List[Any]						=				model(snake_case       ,	attention_mask=snake_case	)
											SCREAMING_SNAKE_CASE							:     Optional[Any]						=				model(snake_case	)
											self.parent.assertEqual(result.last_hidden_state.shape       ,	(self.batch_size, self.seq_length, self.hidden_size)	)







						def       lowerCamelCase_				(  self    :				int       ,	snake_case    :				Optional[int]       ,	snake_case    :				Optional[int]       ,	snake_case    :				str       ,	snake_case    :				Any       ,	snake_case    :				Union[str, Any]       ,	snake_case    :				List[str]       ,	snake_case    :				Optional[Any]       ,	snake_case    :				Tuple       ,	snake_case    :				int       ,	):



											'''simple docstring'''


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







						def       lowerCamelCase_				(  self    :				Tuple       ,	snake_case    :				Optional[Any]       ,	snake_case    :				int       ,	snake_case    :				List[str]       ,	snake_case    :				Dict       ,	snake_case    :				List[str]       ,	snake_case    :				Optional[int]       ,	snake_case    :				Tuple       ,	snake_case    :				List[Any]       ,	snake_case    :				Tuple       ,	):



											'''simple docstring'''


											SCREAMING_SNAKE_CASE							:     str						=				FalconForCausalLM(config=snake_case	)
											model.to(snake_case	)
											model.eval()
											SCREAMING_SNAKE_CASE							:     Any						=				model(snake_case       ,	attention_mask=snake_case       ,	labels=snake_case	)
											self.parent.assertEqual(result.logits.shape       ,	(self.batch_size, self.seq_length, self.vocab_size)	)







						def       lowerCamelCase_				(  self    :				Any       ,	snake_case    :				Optional[Any]       ,	snake_case    :				Any       ,	snake_case    :				Optional[Any]       ,	snake_case    :				Optional[Any]       ,	snake_case    :				Union[str, Any]       ,	snake_case    :				Optional[int]       ,	snake_case    :				Optional[int]       ,	snake_case    :				List[str]       ,	snake_case    :				str       ,	):



											'''simple docstring'''


											SCREAMING_SNAKE_CASE							:     Dict						=				True
											SCREAMING_SNAKE_CASE							:     Union[str, Any]						=				True
											SCREAMING_SNAKE_CASE							:     Any						=				FalconForCausalLM(config=snake_case	)
											model.to(snake_case	)
											model.eval()

											# first forward pass
											SCREAMING_SNAKE_CASE							:     str						=				model(
											    snake_case       ,	attention_mask=snake_case       ,	encoder_hidden_states=snake_case       ,	encoder_attention_mask=snake_case       ,	use_cache=snake_case       ,	)
											SCREAMING_SNAKE_CASE							:     Union[str, Any]						=				outputs.past_key_values

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

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

											SCREAMING_SNAKE_CASE							:     Optional[int]						=				model(
											    snake_case       ,	attention_mask=snake_case       ,	encoder_hidden_states=snake_case       ,	encoder_attention_mask=snake_case       ,	output_hidden_states=snake_case       ,	)['hidden_states'][0]
											SCREAMING_SNAKE_CASE							:     List[str]						=				model(
											    snake_case       ,	attention_mask=snake_case       ,	encoder_hidden_states=snake_case       ,	encoder_attention_mask=snake_case       ,	past_key_values=snake_case       ,	output_hidden_states=snake_case       ,	)['hidden_states'][0]

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







						def       lowerCamelCase_				(  self    :				Tuple	):



											'''simple docstring'''


											SCREAMING_SNAKE_CASE							:     int						=				self.prepare_config_and_inputs()
											(
											    (
											    SCREAMING_SNAKE_CASE
											)					,  (
											    SCREAMING_SNAKE_CASE
											)					,  (
											    SCREAMING_SNAKE_CASE
											)					,  (
											    SCREAMING_SNAKE_CASE
											)					,  (
											    SCREAMING_SNAKE_CASE
											)					,  (
											    SCREAMING_SNAKE_CASE
											)					,  (
											    SCREAMING_SNAKE_CASE
											)					,  
											)							:     str						=				config_and_inputs
											SCREAMING_SNAKE_CASE							:     str						=				{'input_ids': input_ids, 'attention_mask': input_mask}
											return config, inputs_dict


@require_torch
class 						lowercase    (							SCREAMING_SNAKE_CASE_      ,    SCREAMING_SNAKE_CASE_      ,    SCREAMING_SNAKE_CASE_      ,    unittest.TestCase):
						'''simple docstring'''


						UpperCAmelCase			:		Any		=		(
						    (
						        FalconModel,
						        FalconForCausalLM,
						        FalconForSequenceClassification,
						        FalconForTokenClassification,
						        FalconForQuestionAnswering,
						    )
						    if is_torch_available()
						    else ()
						)
						UpperCAmelCase			:		Tuple		=		(FalconForCausalLM,) if is_torch_available() else ()
						UpperCAmelCase			:		Tuple		=		(
						    {
						        'feature-extraction': FalconModel,
						        'text-classification': FalconForSequenceClassification,
						        'text-generation': FalconForCausalLM,
						        'question-answering': FalconForQuestionAnswering,
						        'token-classification': FalconForTokenClassification,
						        'zero-shot': FalconForSequenceClassification,
						    }
						    if is_torch_available()
						    else {}
						)
						UpperCAmelCase			:		Optional[Any]		=		False
						UpperCAmelCase			:		Dict		=		False







						def       lowerCamelCase_				(  self    :				Optional[Any]	):



											'''simple docstring'''


											SCREAMING_SNAKE_CASE							:     int						=				FalconModelTester(self	)
											SCREAMING_SNAKE_CASE							:     Tuple						=				ConfigTester(self       ,	config_class=snake_case       ,	hidden_size=37	)







						def       lowerCamelCase_				(  self    :				Optional[Any]	):



											'''simple docstring'''


											self.config_tester.run_common_tests()







						def       lowerCamelCase_				(  self    :				str	):



											'''simple docstring'''


											SCREAMING_SNAKE_CASE							:     List[Any]						=				self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*snake_case	)







						def       lowerCamelCase_				(  self    :				Any	):



											'''simple docstring'''


											SCREAMING_SNAKE_CASE					,  *SCREAMING_SNAKE_CASE							:     Optional[int]						=				self.model_tester.prepare_config_and_inputs()
											for alibi in [True, False]:
																SCREAMING_SNAKE_CASE							:     Optional[Any]						=				alibi
																self.model_tester.create_and_check_model(snake_case       ,	*snake_case	)







						def       lowerCamelCase_				(  self    :				int	):



											'''simple docstring'''


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







						def       lowerCamelCase_				(  self    :				int	):



											'''simple docstring'''


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







						def       lowerCamelCase_				(  self    :				str	):



											'''simple docstring'''


											SCREAMING_SNAKE_CASE					,  SCREAMING_SNAKE_CASE							:     Union[str, Any]						=				self.model_tester.prepare_config_and_inputs_for_common()
											SCREAMING_SNAKE_CASE							:     List[Any]						=				input_dict['input_ids']
											SCREAMING_SNAKE_CASE							:     Any						=				FalconForCausalLM(snake_case	)
											model.to(snake_case	)
											model.eval()
											SCREAMING_SNAKE_CASE							:     List[str]						=				model(snake_case       ,	use_cache=snake_case	)
											SCREAMING_SNAKE_CASE							:     Optional[int]						=				input_ids.shape[0]
											SCREAMING_SNAKE_CASE							:     Optional[int]						=				model._convert_to_rw_cache(result.past_key_values	)
											SCREAMING_SNAKE_CASE							:     Union[str, Any]						=				model._convert_cache_to_standard_format(snake_case       ,	snake_case	)
											for layer in range(len(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       lowerCamelCase_				(  self    :				str	):



											'''simple docstring'''


											SCREAMING_SNAKE_CASE					,  SCREAMING_SNAKE_CASE							:     Dict						=				self.model_tester.prepare_config_and_inputs_for_common()
											SCREAMING_SNAKE_CASE							:     Tuple						=				3
											SCREAMING_SNAKE_CASE							:     Dict						=				'multi_label_classification'
											SCREAMING_SNAKE_CASE							:     Any						=				input_dict['input_ids']
											SCREAMING_SNAKE_CASE							:     Any						=				input_ids.ne(1	).to(snake_case	)
											SCREAMING_SNAKE_CASE							:     List[str]						=				ids_tensor(
											    [self.model_tester.batch_size, config.num_labels]       ,	self.model_tester.type_sequence_label_size	).to(torch.float	)
											SCREAMING_SNAKE_CASE							:     List[str]						=				FalconForSequenceClassification(snake_case	)
											model.to(snake_case	)
											model.eval()
											SCREAMING_SNAKE_CASE							:     Any						=				model(snake_case       ,	attention_mask=snake_case       ,	labels=snake_case	)
											self.assertEqual(result.logits.shape       ,	(self.model_tester.batch_size, self.model_tester.num_labels)	)







						def       lowerCamelCase_				(  self    :				List[str]	):



											'''simple docstring'''


											for model_class in self.all_generative_model_classes:
																SCREAMING_SNAKE_CASE					,  SCREAMING_SNAKE_CASE							:     Any						=				self.model_tester.prepare_config_and_inputs_for_common()

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

																SCREAMING_SNAKE_CASE							:     Any						=				model_class(snake_case	).to(snake_case	)
																if "use_cache" not in inputs:
																					SCREAMING_SNAKE_CASE							:     Union[str, Any]						=				True
																SCREAMING_SNAKE_CASE							:     Any						=				model(**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

																SCREAMING_SNAKE_CASE							:     Tuple						=				(
																    getattr(snake_case       ,	'decoder_layers'       ,	snake_case	)
																    or getattr(snake_case       ,	'num_decoder_layers'       ,	snake_case	)
																    or config.num_hidden_layers
																)
																SCREAMING_SNAKE_CASE							:     Tuple						=				getattr(snake_case       ,	'num_kv_heads'       ,	config.num_attention_heads	)
																SCREAMING_SNAKE_CASE							:     Optional[int]						=				getattr(snake_case       ,	'd_model'       ,	config.hidden_size	)
																SCREAMING_SNAKE_CASE							:     List[Any]						=				embed_dim // num_attention_heads

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

																SCREAMING_SNAKE_CASE					,  SCREAMING_SNAKE_CASE							:     Tuple						=				inputs['input_ids'].shape
																for i in range(snake_case	):
																					if config.new_decoder_architecture:
																										SCREAMING_SNAKE_CASE							:     Optional[int]						=				config.num_attention_heads
																					elif config.multi_query:
																										SCREAMING_SNAKE_CASE							:     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 						lowercase    (							unittest.TestCase):
						'''simple docstring'''


						@slow
						def       lowerCamelCase_				(  self    :				Tuple	):



											'''simple docstring'''


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

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

											SCREAMING_SNAKE_CASE							:     int						=				model.generate(**snake_case       ,	do_sample=snake_case       ,	max_new_tokens=19	)
											SCREAMING_SNAKE_CASE							:     int						=				tokenizer.batch_decode(snake_case	)[0]

											self.assertEqual(snake_case       ,	snake_case	)







						@slow
						def       lowerCamelCase_				(  self    :				Tuple	):



											'''simple docstring'''


											for repo in ["Rocketknight1/tiny-random-falcon-7b", "Rocketknight1/tiny-random-falcon-40b"]:
																SCREAMING_SNAKE_CASE							:     int						=				AutoTokenizer.from_pretrained(snake_case	)
																SCREAMING_SNAKE_CASE							:     Tuple						=				FalconForCausalLM.from_pretrained(snake_case	)
																model.eval()
																model.to(snake_case	)
																SCREAMING_SNAKE_CASE							:     Tuple						=				tokenizer('My favorite food is'       ,	return_tensors='pt'	).to(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(**snake_case       ,	do_sample=snake_case       ,	max_new_tokens=4	)
																model.generate(**snake_case       ,	do_sample=snake_case       ,	max_new_tokens=4	)
																model.generate(**snake_case       ,	num_beams=2       ,	max_new_tokens=4	)







						@slow
						def       lowerCamelCase_				(  self    :				List[Any]	):



											'''simple docstring'''


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

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




import argparse

import torch
from safetensors.torch import load_file

from diffusers import StableDiffusionPipeline

def        __a			(   __lowerCAmelCase      ,	__lowerCAmelCase      ,	__lowerCAmelCase      ,	__lowerCAmelCase      ,	__lowerCAmelCase   )     ->     str:
					# load base model
					SCREAMING_SNAKE_CASE							:     Dict						=				StableDiffusionPipeline.from_pretrained(__lowerCAmelCase      ,	torch_dtype=torch.floataa   )

					# load LoRA weight from .safetensors
					SCREAMING_SNAKE_CASE							:     Any						=				load_file(__lowerCAmelCase   )

					SCREAMING_SNAKE_CASE							:     Optional[int]						=				[]

					# directly update weight in diffusers model
					for key in state_dict:
										# it is suggested to print out the key, it usually will be something like below
										# "lora_te_text_model_encoder_layers_0_self_attn_k_proj.lora_down.weight"

										# as we have set the alpha beforehand, so just skip
										if ".alpha" in key or key in visited:
															continue

										if "text" in key:
															SCREAMING_SNAKE_CASE							:     Dict						=				key.split('.'   )[0].split(LORA_PREFIX_TEXT_ENCODER + '_'   )[-1].split('_'   )
															SCREAMING_SNAKE_CASE							:     Tuple						=				pipeline.text_encoder
										else:
															SCREAMING_SNAKE_CASE							:     Any						=				key.split('.'   )[0].split(LORA_PREFIX_UNET + '_'   )[-1].split('_'   )
															SCREAMING_SNAKE_CASE							:     Dict						=				pipeline.unet

										# find the target layer
										SCREAMING_SNAKE_CASE							:     Union[str, Any]						=				layer_infos.pop(0   )
										while len(__lowerCAmelCase   ) > -1:
															try:
																				SCREAMING_SNAKE_CASE							:     List[Any]						=				curr_layer.__getattr__(__lowerCAmelCase   )
																				if len(__lowerCAmelCase   ) > 0:
																									SCREAMING_SNAKE_CASE							:     Optional[int]						=				layer_infos.pop(0   )
																				elif len(__lowerCAmelCase   ) == 0:
																									break
															except Exception:
																				if len(__lowerCAmelCase   ) > 0:
																									temp_name += "_" + layer_infos.pop(0   )
																				else:
																									SCREAMING_SNAKE_CASE							:     List[str]						=				layer_infos.pop(0   )

										SCREAMING_SNAKE_CASE							:     str						=				[]
										if "lora_down" in key:
															pair_keys.append(key.replace('lora_down'      ,	'lora_up'   )   )
															pair_keys.append(__lowerCAmelCase   )
										else:
															pair_keys.append(__lowerCAmelCase   )
															pair_keys.append(key.replace('lora_up'      ,	'lora_down'   )   )

										# update weight
										if len(state_dict[pair_keys[0]].shape   ) == 4:
															SCREAMING_SNAKE_CASE							:     Any						=				state_dict[pair_keys[0]].squeeze(3   ).squeeze(2   ).to(torch.floataa   )
															SCREAMING_SNAKE_CASE							:     str						=				state_dict[pair_keys[1]].squeeze(3   ).squeeze(2   ).to(torch.floataa   )
															curr_layer.weight.data += alpha * torch.mm(__lowerCAmelCase      ,	__lowerCAmelCase   ).unsqueeze(2   ).unsqueeze(3   )
										else:
															SCREAMING_SNAKE_CASE							:     List[Any]						=				state_dict[pair_keys[0]].to(torch.floataa   )
															SCREAMING_SNAKE_CASE							:     Optional[int]						=				state_dict[pair_keys[1]].to(torch.floataa   )
															curr_layer.weight.data += alpha * torch.mm(__lowerCAmelCase      ,	__lowerCAmelCase   )

										# update visited list
										for item in pair_keys:
															visited.append(__lowerCAmelCase   )

					return pipeline


if __name__ == "__main__":
			_lowerCamelCase    :		Tuple		      =    argparse.ArgumentParser()

			parser.add_argument(
			    """--base_model_path""", default=None, type=str, required=True, help="""Path to the base model in diffusers format."""
			)
			parser.add_argument(
			    """--checkpoint_path""", default=None, type=str, required=True, help="""Path to the checkpoint to convert."""
			)
			parser.add_argument("""--dump_path""", default=None, type=str, required=True, help="""Path to the output model.""")
			parser.add_argument(
			    """--lora_prefix_unet""", default="""lora_unet""", type=str, help="""The prefix of UNet weight in safetensors"""
			)
			parser.add_argument(
			    """--lora_prefix_text_encoder""",
			    default="""lora_te""",
			    type=str,
			    help="""The prefix of text encoder weight in safetensors""",
			)
			parser.add_argument("""--alpha""", default=0.75, type=float, help="""The merging ratio in W = W0 + alpha * deltaW""")
			parser.add_argument(
			    """--to_safetensors""", action="""store_true""", help="""Whether to store pipeline in safetensors format or not."""
			)
			parser.add_argument("""--device""", type=str, help="""Device to use (e.g. cpu, cuda:0, cuda:1, etc.)""")

			_lowerCamelCase    :		List[Any]		      =    parser.parse_args()

			_lowerCamelCase    :		List[str]		      =    args.base_model_path
			_lowerCamelCase    :		str		      =    args.checkpoint_path
			_lowerCamelCase    :		Dict		      =    args.dump_path
			_lowerCamelCase    :		Any		      =    args.lora_prefix_unet
			_lowerCamelCase    :		List[Any]		      =    args.lora_prefix_text_encoder
			_lowerCamelCase    :		List[Any]		      =    args.alpha

			_lowerCamelCase    :		Optional[Any]		      =    convert(base_model_path, checkpoint_path, lora_prefix_unet, lora_prefix_text_encoder, alpha)

			_lowerCamelCase    :		Optional[Any]		      =    pipe.to(args.device)
			pipe.save_pretrained(args.dump_path, safe_serialization=args.to_safetensors)