




'''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      (  _UpperCamelCase:       Optional[int]=32   ,		_UpperCamelCase:       List[str]=10   ,		_UpperCamelCase:       Any=100   ,		_UpperCamelCase:       Optional[Any]=1_026   ,		_UpperCamelCase:       Union[str, Any]=True   ,		_UpperCamelCase:       Union[str, Any]="data/tokenized_stories_train_wikitext103.jbl"   ,		_UpperCamelCase:       List[Any]="igf_context_pairs.jbl"   ,		)		->					List[Any]:

     """simple docstring"""




     set_seed(3       )
     # generate train_data and objective_set
     _snake_case				,   _snake_case				= generate_datasets(
         _UpperCamelCase   ,		_UpperCamelCase   ,		number=_UpperCamelCase   ,		min_len=1_026   ,		trim=_UpperCamelCase       )
     # 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?
     _snake_case				= torch.device("cuda:0" if torch.cuda.is_available() else "cpu"       )

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

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

     # clean up, delete model and data we don't need anymore
     del model, train_data, objective_set
     torch.cuda.empty_cache()





def       __a      (  _UpperCamelCase:       List[str]   ,		_UpperCamelCase:       str=15   ,		_UpperCamelCase:       List[Any]=128   ,		_UpperCamelCase:       List[Any]=100   ,		_UpperCamelCase:       List[str]="igf_model.pt"   ,		)		->					List[str]:

     """simple docstring"""




     set_seed(42       )

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

     # Initialize secondary learner to use embedding weights of model
     _snake_case				= SecondaryLearner(_UpperCamelCase       )

     # Train secondary learner
     _snake_case				= train_secondary_learner(
         _UpperCamelCase   ,		_UpperCamelCase   ,		max_epochs=_UpperCamelCase   ,		batch_size=_UpperCamelCase   ,		eval_freq=100   ,		igf_model_path=_UpperCamelCase   ,		)

     del model, secondary_learner_train_data
     torch.cuda.empty_cache()

     return secondary_learner





def       __a      (  _UpperCamelCase:       Tuple   ,		_UpperCamelCase:       int   ,		_UpperCamelCase:       Optional[int]   ,		_UpperCamelCase:       Optional[Any]=32   ,		_UpperCamelCase:       Optional[int]=1_000   ,		_UpperCamelCase:       Tuple=16   ,		_UpperCamelCase:       Optional[Any]=1.0   ,		_UpperCamelCase:       Dict=recopy_gpta   ,		_UpperCamelCase:       int=None   ,		_UpperCamelCase:       Dict=10   ,		_UpperCamelCase:       Optional[Any]="gpt2_finetuned.pt"   ,		)		->					Dict:

     """simple docstring"""




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

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

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

     _snake_case				= []
     _snake_case				= []

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

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

                    # 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:
                         _snake_case				= -1
                    if predicted_q < threshold:
                         _snake_case				= False

            # If we passed the filter, add the context to the batch!
               if do_backprop:
                    contexts.append(np.array(context.cpu()       )       )
                    _snake_case				= 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()
                    _snake_case				= 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:
                         _snake_case				= compute_perplexity(_UpperCamelCase   ,		_UpperCamelCase   ,		_UpperCamelCase       )
                         test_perps.append(_UpperCamelCase       )

                         print("Test perplexity, step"   ,		_UpperCamelCase   ,		":"   ,		_UpperCamelCase       )
            # 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()   ,		_UpperCamelCase       )
     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      (  )		->					Optional[int]:

     """simple docstring"""




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

     # Required parameters
     parser.add_argument(
         "--data_dir"   ,		default=_UpperCamelCase   ,		type=_UpperCamelCase   ,		required=_UpperCamelCase   ,		help="The input data dir. Should contain data files for WikiText."   ,		)
     parser.add_argument(
         "--model_name_or_path"   ,		default=_UpperCamelCase   ,		type=_UpperCamelCase   ,		required=_UpperCamelCase   ,		help="Path to pretrained model or model identifier from huggingface.co/models"   ,		)
     parser.add_argument(
         "--data_file"   ,		type=_UpperCamelCase   ,		default=_UpperCamelCase   ,		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=_UpperCamelCase   ,		default=_UpperCamelCase   ,		help="A jbl file containing the context and information gain pairs to train secondary learner."   ,		)

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

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

     parser.add_argument(
         "--context_len"   ,		default=32   ,		type=_UpperCamelCase   ,		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=_UpperCamelCase   ,		help="number of articles that are long enough to be used as our objective set"   ,		)
     parser.add_argument(
         "--eval_freq"   ,		default=100   ,		type=_UpperCamelCase   ,		help="secondary model evaluation is triggered at eval_freq"       )

     parser.add_argument("--max_steps"   ,		default=1_000   ,		type=_UpperCamelCase   ,		help="To calculate training epochs"       )

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

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

     parser.add_argument(
         "--eval_interval"   ,		default=10   ,		type=_UpperCamelCase   ,		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=_UpperCamelCase   ,		help="The number of examples split to be used as objective_set/test_data"       )

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

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

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

     parser.add_argument(
         "--threshold"   ,		default=1.0   ,		type=_UpperCamelCase   ,		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=_UpperCamelCase   ,		help="finetuned_model_name"       )

     parser.add_argument(
         "--recopy_model"   ,		default=_UpperCamelCase   ,		type=_UpperCamelCase   ,		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=1_026   ,		trim=_UpperCamelCase   ,		data_file="data/tokenized_stories_train_wikitext103.jbl"   ,		igf_data_file="igf_context_pairs.jbl"   ,		)

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

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

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

     # Generate train and test data to train and evaluate gpt2 model
     _snake_case				,   _snake_case				= generate_datasets(
         context_len=32   ,		file="data/tokenized_stories_train_wikitext103.jbl"   ,		number=100   ,		min_len=1_026   ,		trim=_UpperCamelCase       )

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


if __name__ == "__main__":
   main()








'''simple docstring'''





from ...configuration_utils import PretrainedConfig
from ...utils import logging


UpperCamelCase_		:     int          =		logging.get_logger(__name__)

UpperCamelCase_		:     Any          =		{
    '''tiiuae/falcon-40b''': '''https://huggingface.co/tiiuae/falcon-40b/resolve/main/config.json''',
    '''tiiuae/falcon-7b''': '''https://huggingface.co/tiiuae/falcon-7b/resolve/main/config.json''',
}

class   _a							(		__lowerCAmelCase				):
     SCREAMING_SNAKE_CASE_				:       List[str]       							=		"""falcon"""
     SCREAMING_SNAKE_CASE_				:       Optional[int]       							=		["""past_key_values"""]






     def __init__(					self   ,_SCREAMING_SNAKE_CASE=65_024   ,_SCREAMING_SNAKE_CASE=4_544   ,_SCREAMING_SNAKE_CASE=32   ,_SCREAMING_SNAKE_CASE=71   ,_SCREAMING_SNAKE_CASE=1e-5   ,_SCREAMING_SNAKE_CASE=0.0_2   ,_SCREAMING_SNAKE_CASE=True   ,_SCREAMING_SNAKE_CASE=0.0   ,_SCREAMING_SNAKE_CASE=0.0   ,_SCREAMING_SNAKE_CASE=None   ,_SCREAMING_SNAKE_CASE=False   ,_SCREAMING_SNAKE_CASE=False   ,_SCREAMING_SNAKE_CASE=True   ,_SCREAMING_SNAKE_CASE=True   ,_SCREAMING_SNAKE_CASE=False   ,_SCREAMING_SNAKE_CASE=11   ,_SCREAMING_SNAKE_CASE=11   ,**_SCREAMING_SNAKE_CASE   ,)				->     Dict:
          _snake_case				= vocab_size
          # Backward compatibility with n_embed kwarg
          _snake_case				= kwargs.pop("n_embed"   ,_SCREAMING_SNAKE_CASE       )
          _snake_case				= hidden_size if n_embed is None else n_embed
          _snake_case				= num_hidden_layers
          _snake_case				= num_attention_heads
          _snake_case				= layer_norm_epsilon
          _snake_case				= initializer_range
          _snake_case				= use_cache
          _snake_case				= hidden_dropout
          _snake_case				= attention_dropout

          _snake_case				= bos_token_id
          _snake_case				= eos_token_id
          _snake_case				= num_attention_heads if num_kv_heads is None else num_kv_heads
          _snake_case				= alibi
          _snake_case				= new_decoder_architecture
          _snake_case				= multi_query  # Ignored when new_decoder_architecture is True
          _snake_case				= parallel_attn
          _snake_case				= bias

          super().__init__(bos_token_id=_SCREAMING_SNAKE_CASE   ,eos_token_id=_SCREAMING_SNAKE_CASE   ,**_SCREAMING_SNAKE_CASE       )






     @property
     def 							_lowercase      (					self       )				->     Union[str, Any]:
          return self.hidden_size // self.num_attention_heads






     @property
     def 							_lowercase      (					self       )				->     Optional[int]:
          return not self.alibi


