
"""simple docstring"""





import argparse
import random

import joblib
import numpy as np
import torch
from igf.igf import (
    SecondaryLearner,
    collect_objective_set,
    compute_perplexity,
    generate_datasets,
    load_gpta,
    recopy_gpta,
    set_seed,
    train_secondary_learner,
)
from torch.utils.data import DataLoader, RandomSampler

from transformers import GPTaLMHeadModel
def 							A_     (__a=32				,   __a=10				,   __a=100				,   __a=1026				,   __a=True				,   __a="data/tokenized_stories_train_wikitext103.jbl"				,   __a="igf_context_pairs.jbl"				,   ):

					'''simple docstring'''





					set_seed(3				)
					# generate train_data and objective_set
					A_				 =						generate_datasets(
					    _lowerCamelCase				,   _lowerCamelCase				,   number=_lowerCamelCase				,   min_len=1026				,   trim=_lowerCamelCase				)
					# keeps model same across runs
					set_seed(4				)
					# model, lm_optimizer, lm_scheduler = recopy_gpt2(model, device, max_steps) # store original model weights
					# can we train on GPU?
					A_				 =						torch.device("cuda:0" if torch.cuda.is_available() else "cpu"				)

					# load pretrained model
					A_				 =						load_gpta("gpt2"				).to(_lowerCamelCase				)
					print("computing perplexity on objective set"				)
					A_				 =						compute_perplexity(_lowerCamelCase				,   _lowerCamelCase				,   _lowerCamelCase				).item()
					print("perplexity on objective set:"				,   _lowerCamelCase				)

					# collect igf pairs and save to file demo.jbl
					collect_objective_set(_lowerCamelCase				,   _lowerCamelCase				,   _lowerCamelCase				,   _lowerCamelCase				,   _lowerCamelCase				,   _lowerCamelCase				,   _lowerCamelCase				,   _lowerCamelCase				)

					# clean up, delete model and data we don't need anymore
					del model, train_data, objective_set
					torch.cuda.empty_cache()
def 							A_     (__a				,   __a=15				,   __a=128				,   __a=100				,   __a="igf_model.pt"				,   ):

					'''simple docstring'''





					set_seed(42				)

					# Load pre-trained model
					A_				 =						GPTaLMHeadModel.from_pretrained("gpt2"				)

					# Initialize secondary learner to use embedding weights of model
					A_				 =						SecondaryLearner(_lowerCamelCase				)

					# Train secondary learner
					A_				 =						train_secondary_learner(
					    _lowerCamelCase				,   _lowerCamelCase				,   max_epochs=_lowerCamelCase				,   batch_size=_lowerCamelCase				,   eval_freq=100				,   igf_model_path=_lowerCamelCase				,   )

					del model, secondary_learner_train_data
					torch.cuda.empty_cache()

					return secondary_learner
def 							A_     (__a				,   __a				,   __a				,   __a=32				,   __a=1000				,   __a=16				,   __a=1.0				,   __a=recopy_gpta				,   __a=None				,   __a=10				,   __a="gpt2_finetuned.pt"				,   ):

					'''simple docstring'''





					A_				 =						torch.device("cuda:0" if torch.cuda.is_available() else "cpu"				)
					A_				 =						RandomSampler(_lowerCamelCase				)
					A_				 =						DataLoader(_lowerCamelCase				,   sampler=_lowerCamelCase				)

					A_				 =						max_steps // (len(_lowerCamelCase				)) + 1
					A_				 =						0
					A_				 =						torch.zeros((1, context_len)				,   dtype=torch.long				,   device=_lowerCamelCase				)
					A_				 =						recopy_model(_lowerCamelCase				,   _lowerCamelCase				,   _lowerCamelCase				)

					model.train()
					if secondary_learner is not None:
										secondary_learner.to(_lowerCamelCase				)
										secondary_learner.eval()
					A_				 =						[]
					A_				 =						0

					A_				 =						[]
					A_				 =						[]

					# Compute the performance of the transformer model at the beginning
					A_				 =						compute_perplexity(_lowerCamelCase				,   _lowerCamelCase				,   _lowerCamelCase				)
					test_perps.append(_lowerCamelCase				)
					print("Test perplexity, step"				,   _lowerCamelCase				,   ":"				,   _lowerCamelCase				)
					for epoch in range(int(_lowerCamelCase				)				):
										for step, example in enumerate(_lowerCamelCase				):
															torch.cuda.empty_cache()
															A_				 =						random.randint(0				,   example.size(2				) - context_len - 1				)
															A_				 =						example[0, 0, start : start + context_len]
															lm_optimizer.zero_grad()
															A_				 =						model(_lowerCamelCase				,   labels=_lowerCamelCase				)
															A_				 =						True

															if secondary_learner is not None:
																				A_				 =						secondary_learner.forward(
																				    torch.tensor(_lowerCamelCase				,   dtype=torch.long				,   device=_lowerCamelCase				).unsqueeze(0				)				)[0].item()
																				observed_qs.append(float(_lowerCamelCase				)				)

																				# Here we implement the simple non-constant threshold for the predicted IG(X) value
																				# We will decay the selectivity of our secondary learner filter from
																				# 1 standard deviation above average to 1 below average after 10 batches.

																				if global_step == 10:
																									A_				 =						-1
																				if predicted_q < threshold:
																									A_				 =						False

            # If we passed the filter, add the context to the batch!
															if do_backprop:
																				contexts.append(np.array(context.cpu()				)				)
																				A_				 =						outputs[0]
																				lm_loss.backward()
																				examples += 1

															del outputs

															# Once the batch is filled with enough contexts, backprop on the batch.
															if examples == batch_size:
																				torch.cuda.empty_cache()
																				A_				 =						0
																				# Do LM backprop
																				torch.nn.utils.clip_grad_norm_(model.parameters()				,   3.0				)
																				lm_optimizer.step()
																				lm_scheduler.step()  # Update learning rate schedule
																				global_step += 1
																				# Compute the performance of the transformer model at this batch
																				if global_step % eval_interval == 0:
																									A_				 =						compute_perplexity(_lowerCamelCase				,   _lowerCamelCase				,   _lowerCamelCase				)
																									test_perps.append(_lowerCamelCase				)

																									print("Test perplexity, step"				,   _lowerCamelCase				,   ":"				,   _lowerCamelCase				)
            # Break out of the loop after 60 batches
															if max_steps > 0 and global_step > 60:
																				break
										if max_steps > 0 and global_step > 60:
															break

    # save finetuned transformer model
					torch.save(model.state_dict()				,   _lowerCamelCase				)
					torch.cuda.empty_cache()
					# Do some cleaning up so we can reinitialize for the next run of this function
					del lm_optimizer
					del lm_scheduler
					return model
def 							A_     ():

					'''simple docstring'''





					A_				 =						argparse.ArgumentParser(description="Fine-tune a transformer model with IGF on a language modeling task"				)

					# Required parameters
					parser.add_argument(
					    "--data_dir"				,   default=_lowerCamelCase				,   type=_lowerCamelCase				,   required=_lowerCamelCase				,   help="The input data dir. Should contain data files for WikiText."				,   )
					parser.add_argument(
					    "--model_name_or_path"				,   default=_lowerCamelCase				,   type=_lowerCamelCase				,   required=_lowerCamelCase				,   help="Path to pretrained model or model identifier from huggingface.co/models"				,   )
					parser.add_argument(
					    "--data_file"				,   type=_lowerCamelCase				,   default=_lowerCamelCase				,   help=(
					        "A jbl file containing tokenized data which can be split as objective dataset, "
					        "train_dataset and test_dataset."
					    )				,   )

					parser.add_argument(
					    "--igf_data_file"				,   type=_lowerCamelCase				,   default=_lowerCamelCase				,   help="A jbl file containing the context and information gain pairs to train secondary learner."				,   )

					parser.add_argument(
					    "--output_dir"				,   default=_lowerCamelCase				,   type=_lowerCamelCase				,   required=_lowerCamelCase				,   help="The output directory where the final fine-tuned model is stored."				,   )

					parser.add_argument(
					    "--tokenizer_name"				,   default=_lowerCamelCase				,   type=_lowerCamelCase				,   help="Pretrained tokenizer name or path if not the same as model_name"				,   )
					parser.add_argument("--seed"				,   type=_lowerCamelCase				,   default=_lowerCamelCase				,   help="A seed for reproducible training."				)

					parser.add_argument(
					    "--context_len"				,   default=32				,   type=_lowerCamelCase				,   help=(
					        "The maximum total input sequence length after tokenization. Sequences longer "
					        "than this will be truncated, sequences shorter will be padded."
					    )				,   )

					parser.add_argument(
					    "--size_objective_set"				,   default=100				,   type=_lowerCamelCase				,   help="number of articles that are long enough to be used as our objective set"				,   )
					parser.add_argument(
					    "--eval_freq"				,   default=100				,   type=_lowerCamelCase				,   help="secondary model evaluation is triggered at eval_freq"				)

					parser.add_argument("--max_steps"				,   default=1000				,   type=_lowerCamelCase				,   help="To calculate training epochs"				)

					parser.add_argument(
					    "--secondary_learner_batch_size"				,   default=128				,   type=_lowerCamelCase				,   help="batch size of training data for secondary learner"				,   )

					parser.add_argument(
					    "--batch_size"				,   default=16				,   type=_lowerCamelCase				,   help="batch size of training data of language model(gpt2) "				)

					parser.add_argument(
					    "--eval_interval"				,   default=10				,   type=_lowerCamelCase				,   help=(
					        "decay the selectivity of our secondary learner filter from"
					        "1 standard deviation above average to 1 below average after 10 batches"
					    )				,   )

					parser.add_argument(
					    "--number"				,   default=100				,   type=_lowerCamelCase				,   help="The number of examples split to be used as objective_set/test_data"				)

					parser.add_argument(
					    "--min_len"				,   default=1026				,   type=_lowerCamelCase				,   help="The minimum length of the article to be used as objective set"				)

					parser.add_argument(
					    "--secondary_learner_max_epochs"				,   default=15				,   type=_lowerCamelCase				,   help="number of epochs to train secondary learner"				)

					parser.add_argument("--trim"				,   default=_lowerCamelCase				,   type=_lowerCamelCase				,   help="truncate the example if it exceeds context length"				)

					parser.add_argument(
					    "--threshold"				,   default=1.0				,   type=_lowerCamelCase				,   help=(
					        "The threshold value used by secondary learner to filter the train_data and allow only"
					        " informative data as input to the model"
					    )				,   )

					parser.add_argument("--finetuned_model_name"				,   default="gpt2_finetuned.pt"				,   type=_lowerCamelCase				,   help="finetuned_model_name"				)

					parser.add_argument(
					    "--recopy_model"				,   default=_lowerCamelCase				,   type=_lowerCamelCase				,   help="Reset the model to the original pretrained GPT-2 weights after each iteration"				,   )

					# function calls
					# Collecting *n* pairs of context and information gain(X, IG(X)) for training the secondary learner
					generate_n_pairs(
					    context_len=32				,   max_steps=10				,   size_objective_set=100				,   min_len=1026				,   trim=_lowerCamelCase				,   data_file="data/tokenized_stories_train_wikitext103.jbl"				,   igf_data_file="igf_context_pairs.jbl"				,   )

					# Load train data for secondary learner
					A_				 =						joblib.load("data/IGF_values.jbl"				)

					# Train secondary learner
					A_				 =						training_secondary_learner(
					    _lowerCamelCase				,   secondary_learner_max_epochs=15				,   secondary_learner_batch_size=128				,   eval_freq=100				,   igf_model_path="igf_model.pt"				,   )

					# load pretrained gpt2 model
					A_				 =						GPTaLMHeadModel.from_pretrained("gpt2"				)
					set_seed(42				)

					# Generate train and test data to train and evaluate gpt2 model
					A_				 =						generate_datasets(
					    context_len=32				,   file="data/tokenized_stories_train_wikitext103.jbl"				,   number=100				,   min_len=1026				,   trim=_lowerCamelCase				)

					# fine-tuning of the gpt2 model using igf (Information Gain Filtration)
					finetune(
					    _lowerCamelCase				,   _lowerCamelCase				,   _lowerCamelCase				,   context_len=32				,   max_steps=1000				,   batch_size=16				,   threshold=1.0				,   recopy_model=_lowerCamelCase				,   secondary_learner=_lowerCamelCase				,   eval_interval=10				,   finetuned_model_name="gpt2_finetuned.pt"				,   )


if __name__ == "__main__":
		main()







"""simple docstring"""





import inspect
import unittest

from huggingface_hub import hf_hub_download

from transformers import ASTConfig
from transformers.testing_utils import require_torch, require_torchaudio, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_torchaudio_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
		import torch
		from torch import nn

		from transformers import ASTForAudioClassification, ASTModel
		from transformers.models.audio_spectrogram_transformer.modeling_audio_spectrogram_transformer import (
		    AUDIO_SPECTROGRAM_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
		)


if is_torchaudio_available():
		import torchaudio

		from transformers import ASTFeatureExtractor



class __lowerCAmelCase				:




			"""simple docstring"""







			def __init__(				self						:					Union[str, Any]       ,   _snake_case						:					Optional[int]       ,   _snake_case						:					Any=13       ,   _snake_case						:					Union[str, Any]=2       ,   _snake_case						:					Optional[int]=24       ,   _snake_case						:					Optional[Any]=16       ,   _snake_case						:					List[str]=True       ,   _snake_case						:					str=True       ,   _snake_case						:					List[Any]=32       ,   _snake_case						:					str=5       ,   _snake_case						:					int=4       ,   _snake_case						:					List[str]=37       ,   _snake_case						:					int="gelu"       ,   _snake_case						:					str=0.1       ,   _snake_case						:					Optional[int]=0.1       ,   _snake_case						:					Optional[int]=10       ,   _snake_case						:					int=0.0_2       ,   _snake_case						:					int=None       ,   _snake_case						:					Optional[Any]=2       ,   _snake_case						:					int=2       ,   )      -> Dict:







								"""simple docstring"""
								A_				 =						parent
								A_				 =						batch_size
								A_				 =						patch_size
								A_				 =						max_length
								A_				 =						num_mel_bins
								A_				 =						is_training
								A_				 =						use_labels
								A_				 =						hidden_size
								A_				 =						num_hidden_layers
								A_				 =						num_attention_heads
								A_				 =						intermediate_size
								A_				 =						hidden_act
								A_				 =						hidden_dropout_prob
								A_				 =						attention_probs_dropout_prob
								A_				 =						type_sequence_label_size
								A_				 =						initializer_range
								A_				 =						scope
								A_				 =						frequency_stride
								A_				 =						time_stride

								# in AST, the seq length equals the number of patches + 2 (we add 2 for the [CLS] and distillation tokens)
								A_				 =						(self.num_mel_bins - self.patch_size) // self.frequency_stride + 1
								A_				 =						(self.max_length - self.patch_size) // self.time_stride + 1
								A_				 =						frequency_out_dimension * time_out_dimension
								A_				 =						num_patches + 2




			def      lowerCamelCase__     (				self						:					Dict	)      -> int:







								"""simple docstring"""
								A_				 =						floats_tensor([self.batch_size, self.max_length, self.num_mel_bins]	)

								A_				 =						None
								if self.use_labels:
													A_				 =						ids_tensor([self.batch_size]       ,   self.type_sequence_label_size	)

								A_				 =						self.get_config()

								return config, input_values, labels




			def      lowerCamelCase__     (				self						:					List[str]	)      -> str:







								"""simple docstring"""
								return ASTConfig(
								    patch_size=self.patch_size       ,   max_length=self.max_length       ,   num_mel_bins=self.num_mel_bins       ,   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       ,   is_decoder=_snake_case       ,   initializer_range=self.initializer_range       ,   frequency_stride=self.frequency_stride       ,   time_stride=self.time_stride       ,   )




			def      lowerCamelCase__     (				self						:					int       ,   _snake_case						:					Dict       ,   _snake_case						:					Tuple       ,   _snake_case						:					Optional[int]	)      -> Union[str, Any]:







								"""simple docstring"""
								A_				 =						ASTModel(config=_snake_case	)
								model.to(_snake_case	)
								model.eval()
								A_				 =						model(_snake_case	)
								self.parent.assertEqual(result.last_hidden_state.shape       ,   (self.batch_size, self.seq_length, self.hidden_size)	)






			def      lowerCamelCase__     (				self						:					str	)      -> List[str]:







								"""simple docstring"""
								A_				 =						self.prepare_config_and_inputs()
								(
								    (
								    A_
								) ,    (
								    A_
								) ,    (
								    A_
								) ,    
								)				 =						config_and_inputs
								A_				 =						{"input_values": input_values}
								return config, inputs_dict





@require_torch
class __lowerCAmelCase				(					_lowercase     ,				_lowercase     ,				unittest.TestCase      ):




			"""simple docstring"""







			snake_case	      =				(
			    (
			        ASTModel,
			        ASTForAudioClassification,
			    )
			    if is_torch_available()
			    else ()
			)
			snake_case	      =				(
			    {"audio-classification": ASTForAudioClassification, "feature-extraction": ASTModel}
			    if is_torch_available()
			    else {}
			)
			snake_case	      =				False
			snake_case	      =				False
			snake_case	      =				False
			snake_case	      =				False




			def      lowerCamelCase__     (				self						:					Any       ,   _snake_case						:					Tuple       ,   _snake_case						:					str       ,   _snake_case						:					List[str]       ,   _snake_case						:					int       ,   _snake_case						:					int	)      -> str:







								"""simple docstring"""
								if pipeline_test_casse_name == "AudioClassificationPipelineTests":
													return True

								return False




			def      lowerCamelCase__     (				self						:					str	)      -> str:







								"""simple docstring"""
								A_				 =						ASTModelTester(self	)
								A_				 =						ConfigTester(self       ,   config_class=_snake_case       ,   has_text_modality=_snake_case       ,   hidden_size=37	)




			def      lowerCamelCase__     (				self						:					List[str]	)      -> str:







								"""simple docstring"""
								self.config_tester.run_common_tests()




			@unittest.skip(reason="AST does not use inputs_embeds"	)
			def      lowerCamelCase__     (				self						:					Union[str, Any]	)      -> Dict:







								"""simple docstring"""
								pass




			def      lowerCamelCase__     (				self						:					List[str]	)      -> str:







								"""simple docstring"""
								A_ ,    A_				 =						self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
													A_				 =						model_class(_snake_case	)
													self.assertIsInstance(model.get_input_embeddings()       ,   (nn.Module)	)
													A_				 =						model.get_output_embeddings()
													self.assertTrue(x is None or isinstance(_snake_case       ,   nn.Linear	)	)




			def      lowerCamelCase__     (				self						:					Any	)      -> Tuple:







								"""simple docstring"""
								A_ ,    A_				 =						self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
													A_				 =						model_class(_snake_case	)
													A_				 =						inspect.signature(model.forward	)
													# signature.parameters is an OrderedDict => so arg_names order is deterministic
													A_				 =						[*signature.parameters.keys()]

													A_				 =						["input_values"]
													self.assertListEqual(arg_names[:1]       ,   _snake_case	)




			def      lowerCamelCase__     (				self						:					Tuple	)      -> Optional[int]:







								"""simple docstring"""
								A_				 =						self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*_snake_case	)






			@slow
			def      lowerCamelCase__     (				self						:					List[Any]	)      -> Tuple:







								"""simple docstring"""
								for model_name in AUDIO_SPECTROGRAM_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													A_				 =						ASTModel.from_pretrained(_snake_case	)
													self.assertIsNotNone(_snake_case	)





def 							A_     ():

					'''simple docstring'''





					A_				 =						hf_hub_download(
					    repo_id="nielsr/audio-spectogram-transformer-checkpoint"				,   filename="sample_audio.flac"				,   repo_type="dataset"				)

					A_ ,    A_				 =						torchaudio.load(__a				)

					return audio, sampling_rate




@require_torch
@require_torchaudio
class __lowerCAmelCase				(					unittest.TestCase      ):




			"""simple docstring"""







			@cached_property
			def      lowerCamelCase__     (				self						:					List[str]	)      -> int:







								"""simple docstring"""
								return (
								    ASTFeatureExtractor.from_pretrained("MIT/ast-finetuned-audioset-10-10-0.4593"	)
								    if is_torchaudio_available()
								    else None
								)






			@slow
			def      lowerCamelCase__     (				self						:					Any	)      -> Optional[int]:







								"""simple docstring"""
								A_				 =						self.default_feature_extractor
								A_				 =						ASTForAudioClassification.from_pretrained("MIT/ast-finetuned-audioset-10-10-0.4593"	).to(_snake_case	)

								A_				 =						self.default_feature_extractor
								A_ ,    A_				 =						prepare_audio()
								A_				 =						audio.squeeze().numpy()
								A_				 =						feature_extractor(_snake_case       ,   sampling_rate=_snake_case       ,   return_tensors="pt"	).to(_snake_case	)

								# forward pass
								with torch.no_grad():
													A_				 =						model(**_snake_case	)

								# verify the logits
								A_				 =						torch.Size((1, 527)	)
								self.assertEqual(outputs.logits.shape       ,   _snake_case	)

								A_				 =						torch.tensor([-0.8_7_6_0, -7.0_0_4_2, -8.6_6_0_2]	).to(_snake_case	)

								self.assertTrue(torch.allclose(outputs.logits[0, :3]       ,   _snake_case       ,   atol=1e-4	)	)




