


"""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     __lowerCAmelCase		(		lowercase    :    int=32		,				lowercase    :    Dict=10		,				lowercase    :    Dict=100		,				lowercase    :    str=1026		,				lowercase    :    Any=True		,				lowercase    :    Dict="data/tokenized_stories_train_wikitext103.jbl"		,				lowercase    :    Union[str, Any]="igf_context_pairs.jbl"		,				)      ->		int:


    """simple docstring"""





    set_seed(3					)
    # generate train_data and objective_set
    snake_case   ,snake_case       : int    =       generate_datasets(
        lowercase		,				lowercase		,				number=lowercase		,				min_len=1026		,				trim=lowercase					)
    # 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       : Optional[int]    =       torch.device("cuda:0" if torch.cuda.is_available() else "cpu"					)

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

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

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







def     __lowerCAmelCase		(		lowercase    :    str		,				lowercase    :    Optional[Any]=15		,				lowercase    :    Any=128		,				lowercase    :    Optional[int]=100		,				lowercase    :    Any="igf_model.pt"		,				)      ->		List[Any]:


    """simple docstring"""





    set_seed(42					)

    # Load pre-trained model
    snake_case       : int    =       GPTaLMHeadModel.from_pretrained("gpt2"					)

    # Initialize secondary learner to use embedding weights of model
    snake_case       : str    =       SecondaryLearner(lowercase					)

    # Train secondary learner
    snake_case       : Optional[Any]    =       train_secondary_learner(
        lowercase		,				lowercase		,				max_epochs=lowercase		,				batch_size=lowercase		,				eval_freq=100		,				igf_model_path=lowercase		,				)

    del model, secondary_learner_train_data
    torch.cuda.empty_cache()

    return secondary_learner







def     __lowerCAmelCase		(		lowercase    :    Dict		,				lowercase    :    Any		,				lowercase    :    str		,				lowercase    :    List[Any]=32		,				lowercase    :    List[Any]=1000		,				lowercase    :    str=16		,				lowercase    :    Optional[int]=1.0		,				lowercase    :    Dict=recopy_gpta		,				lowercase    :    str=None		,				lowercase    :    Tuple=10		,				lowercase    :    Any="gpt2_finetuned.pt"		,				)      ->		Union[str, Any]:


    """simple docstring"""





    snake_case       : Optional[Any]    =       torch.device("cuda:0" if torch.cuda.is_available() else "cpu"					)
    snake_case       : List[str]    =       RandomSampler(lowercase					)
    snake_case       : List[str]    =       DataLoader(lowercase		,				sampler=lowercase					)

    snake_case       : Union[str, Any]    =       max_steps // (len(lowercase					)) + 1
    snake_case       : Dict    =       0
    snake_case       : List[Any]    =       torch.zeros((1, context_len)		,				dtype=torch.long		,				device=lowercase					)
    snake_case   ,snake_case   ,snake_case       : List[str]    =       recopy_model(lowercase		,				lowercase		,				lowercase					)

    model.train()
    if secondary_learner is not None:
        secondary_learner.to(lowercase					)
        secondary_learner.eval()
    snake_case       : List[Any]    =       []
    snake_case       : Dict    =       0

    snake_case       : int    =       []
    snake_case       : Optional[Any]    =       []

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

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

                # 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       : Tuple    =       -1
                if predicted_q < threshold:
                    snake_case       : List[Any]    =       False

            # If we passed the filter, add the context to the batch!
            if do_backprop:
                contexts.append(np.array(context.cpu()					)					)
                snake_case       : Union[str, Any]    =       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       : str    =       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       : Tuple    =       compute_perplexity(lowercase		,				lowercase		,				lowercase					)
                    test_perps.append(lowercase					)

                    print("Test perplexity, step"		,				lowercase		,				":"		,				lowercase					)
            # 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()		,				lowercase					)
    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     __lowerCAmelCase		(		)      ->		List[str]:


    """simple docstring"""





    snake_case       : List[Any]    =       argparse.ArgumentParser(description="Fine-tune a transformer model with IGF on a language modeling task"					)

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

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

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

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

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

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

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

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

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

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

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

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

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

    # Load train data for secondary learner
    snake_case       : List[str]    =       joblib.load("data/IGF_values.jbl"					)

    # Train secondary learner
    snake_case       : Optional[Any]    =       training_secondary_learner(
        lowercase		,				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       : int    =       GPTaLMHeadModel.from_pretrained("gpt2"					)
    set_seed(42					)

    # Generate train and test data to train and evaluate gpt2 model
    snake_case   ,snake_case       : Optional[int]    =       generate_datasets(
        context_len=32		,				file="data/tokenized_stories_train_wikitext103.jbl"		,				number=100		,				min_len=1026		,				trim=lowercase					)

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


if __name__ == "__main__":
  main()









"""simple docstring"""



import uuid
from typing import Any, Dict, List, Optional, Union

from ..utils import add_end_docstrings, is_tf_available, is_torch_available, logging
from .base import PIPELINE_INIT_ARGS, Pipeline


if is_tf_available():
  import tensorflow as tf

if is_torch_available():
  import torch


__snake_case       =       logging.get_logger(__name__)







class _lowerCAmelCase							:







   def __init__(      self		,      UpperCamelCase__ = None		,      UpperCamelCase__ = None		,      UpperCamelCase__=None		,      UpperCamelCase__=None		)      ->						Optional[Any]:

       '''simple docstring'''




       if not conversation_id:
           snake_case       : Union[str, Any]    =       uuid.uuida()
       if past_user_inputs is None:
           snake_case       : Optional[Any]    =       []
       if generated_responses is None:
           snake_case       : Optional[Any]    =       []

       snake_case       : uuid.UUID    =       conversation_id
       snake_case       : List[str]    =       past_user_inputs
       snake_case       : List[str]    =       generated_responses
       snake_case       : Optional[str]    =       text







   def __eq__(      self		,      UpperCamelCase__		)      ->						Any:

       '''simple docstring'''




       if not isinstance(UpperCamelCase__		,      UpperCamelCase__		):
           return False
       if self.uuid == other.uuid:
           return True
       return (
           self.new_user_input == other.new_user_input
           and self.past_user_inputs == other.past_user_inputs
           and self.generated_responses == other.generated_responses
       )







   def 					lowerCamelCase						(      self		,      UpperCamelCase__		,      UpperCamelCase__ = False		)      ->						Dict:

       '''simple docstring'''




       if self.new_user_input:
           if overwrite:
               logger.warning(
                   F'User input added while unprocessed input was existing: "{self.new_user_input}" was overwritten '
                   F'with: "{text}".'		)
               snake_case       : int    =       text
           else:
               logger.warning(
                   F'User input added while unprocessed input was existing: "{self.new_user_input}" new input '
                   F'ignored: "{text}". Set `overwrite` to True to overwrite unprocessed user input'		)
       else:
           snake_case       : Any    =       text







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

       '''simple docstring'''




       if self.new_user_input:
           self.past_user_inputs.append(self.new_user_input		)
       snake_case       : Dict    =       None







   def 					lowerCamelCase						(      self		,      UpperCamelCase__		)      ->						List[str]:

       '''simple docstring'''




       self.generated_responses.append(UpperCamelCase__		)







   def 					lowerCamelCase						(      self		)      ->						List[Any]:

       '''simple docstring'''




       for user_input, generated_response in zip(self.past_user_inputs		,      self.generated_responses		):
           yield True, user_input
           yield False, generated_response
       if self.new_user_input:
           yield True, self.new_user_input







   def __repr__(      self		)      ->						Union[str, Any]:

       '''simple docstring'''




       snake_case       : Optional[Any]    =       F'Conversation id: {self.uuid} \n'
       for is_user, text in self.iter_texts():
           snake_case       : List[str]    =       "user" if is_user else "bot"
           output += F'{name} >> {text} \n'
       return output







@add_end_docstrings(
    snake_case_       ,				R'''
        min_length_for_response (`int`, *optional*, defaults to 32):
            The minimum length (in number of tokens) for a response.
        minimum_tokens (`int`, *optional*, defaults to 10):
            The minimum length of tokens to leave for a response.
    '''       ,				)
class _lowerCAmelCase							( snake_case_      ):







   def __init__(      self		,      *UpperCamelCase__		,      **UpperCamelCase__		)      ->						str:

       '''simple docstring'''




       super().__init__(*UpperCamelCase__		,      **UpperCamelCase__		)
       if self.tokenizer.pad_token_id is None:
           snake_case       : str    =       self.tokenizer.eos_token







   def 					lowerCamelCase						(      self		,      UpperCamelCase__=None		,      UpperCamelCase__=None		,      UpperCamelCase__=None		,      **UpperCamelCase__		)      ->						Union[str, Any]:

       '''simple docstring'''




       snake_case       : Tuple    =       {}
       snake_case       : Optional[int]    =       {}
       snake_case       : Optional[Any]    =       {}

       if min_length_for_response is not None:
           snake_case       : int    =       min_length_for_response
       if minimum_tokens is not None:
           snake_case       : Dict    =       minimum_tokens

       if "max_length" in generate_kwargs:
           snake_case       : Any    =       generate_kwargs["max_length"]
           # self.max_length = generate_kwargs.get("max_length", self.model.config.max_length)
       if clean_up_tokenization_spaces is not None:
           snake_case       : Optional[Any]    =       clean_up_tokenization_spaces

       if generate_kwargs:
           forward_params.update(UpperCamelCase__		)
       return preprocess_params, forward_params, postprocess_params







   def __call__(      self		,      UpperCamelCase__		,      UpperCamelCase__=0		,      **UpperCamelCase__		)      ->						List[str]:

       '''simple docstring'''




       snake_case       : int    =       super().__call__(UpperCamelCase__		,      num_workers=UpperCamelCase__		,      **UpperCamelCase__		)
       if isinstance(UpperCamelCase__		,      UpperCamelCase__		) and len(UpperCamelCase__		) == 1:
           return outputs[0]
       return outputs







   def 					lowerCamelCase						(      self		,      UpperCamelCase__		,      UpperCamelCase__=32		)      ->						Dict[str, Any]:

       '''simple docstring'''




       if not isinstance(UpperCamelCase__		,      UpperCamelCase__		):
           raise ValueError("ConversationalPipeline, expects Conversation as inputs"		)
       if conversation.new_user_input is None:
           raise ValueError(
               F'Conversation with UUID {type(conversation.uuid		)} does not contain new user input to process. '
               "Add user inputs with the conversation's `add_user_input` method"		)
       if hasattr(self.tokenizer		,      "_build_conversation_input_ids"		):
           snake_case       : Optional[Any]    =       self.tokenizer._build_conversation_input_ids(UpperCamelCase__		)
       else:
           # If the tokenizer cannot handle conversations, we default to only the old version
           snake_case       : Optional[Any]    =       self._legacy_parse_and_tokenize(UpperCamelCase__		)

       if self.framework == "pt":
           snake_case       : Union[str, Any]    =       torch.LongTensor([input_ids]		)
       elif self.framework == "tf":
           snake_case       : Optional[int]    =       tf.constant([input_ids]		)
       return {"input_ids": input_ids, "conversation": conversation}







   def 					lowerCamelCase						(      self		,      UpperCamelCase__		,      UpperCamelCase__=10		,      **UpperCamelCase__		)      ->						Union[str, Any]:

       '''simple docstring'''




       snake_case       : Dict    =       generate_kwargs.get("max_length"		,      self.model.config.max_length		)

       snake_case       : Optional[Any]    =       model_inputs["input_ids"].shape[1]
       if max_length - minimum_tokens < n:
           logger.warning(F'Conversation input is to long ({n}), trimming it to ({max_length} - {minimum_tokens})'		)
           snake_case       : List[str]    =       max_length - minimum_tokens
           snake_case       : Dict    =       model_inputs["input_ids"][:, -trim:]
           if "attention_mask" in model_inputs:
               snake_case       : List[Any]    =       model_inputs["attention_mask"][:, -trim:]
       snake_case       : Union[str, Any]    =       model_inputs.pop("conversation"		)
       snake_case       : Union[str, Any]    =       max_length
       snake_case       : Any    =       self.model.generate(**UpperCamelCase__		,      **UpperCamelCase__		)
       if self.model.config.is_encoder_decoder:
           snake_case       : Optional[int]    =       1
       else:
           snake_case       : int    =       n
       return {"output_ids": output_ids[:, start_position:], "conversation": conversation}







   def 					lowerCamelCase						(      self		,      UpperCamelCase__		,      UpperCamelCase__=True		)      ->						Tuple:

       '''simple docstring'''




       snake_case       : Union[str, Any]    =       model_outputs["output_ids"]
       snake_case       : str    =       self.tokenizer.decode(
           output_ids[0]		,      skip_special_tokens=UpperCamelCase__		,      clean_up_tokenization_spaces=UpperCamelCase__		,      )
       snake_case       : List[Any]    =       model_outputs["conversation"]
       conversation.mark_processed()
       conversation.append_response(UpperCamelCase__		)
       return conversation







   def 					lowerCamelCase						(      self		,      UpperCamelCase__		)      ->						Dict:

       '''simple docstring'''




       snake_case       : str    =       self.tokenizer.eos_token_id
       snake_case       : str    =       []
       for is_user, text in conversation.iter_texts():
           if eos_token_id is not None:
               input_ids.extend(self.tokenizer.encode(UpperCamelCase__		,      add_special_tokens=UpperCamelCase__		) + [eos_token_id]		)
           else:
               input_ids.extend(self.tokenizer.encode(UpperCamelCase__		,      add_special_tokens=UpperCamelCase__		)		)

       if len(UpperCamelCase__		) > self.tokenizer.model_max_length:
           snake_case       : str    =       input_ids[-self.tokenizer.model_max_length :]
       return input_ids





