docstr: Clean the datasets by removing too long examples
  The Reward Dataset constraints are:
  - user_input + completion < Reward model max sequence length
  The Actor Dataset constraints are:
  - user_input + completion < Actor model max sequence length
  The RLHF Training Dataset constraints are:
  - user_input + min_completion < Actor model max sequence length
  - user_input + min_completion < Critic model max sequence length
  - user_input + min_completion < Reward model max sequence length
  
  Args:
      config (Config): config object
docstr_tok: ["Clean","the","datasets","by","removing","too","long","examples","The","Reward","Dataset","constraints","are",":","-","user_input","+","completion","<","Reward","model","max","sequence","length","The","Actor","Dataset","constraints","are",":","-","user_input","+","completion","<","Actor","model","max","sequence","length","The","RLHF","Training","Dataset","constraints","are",":","-","user_input","+","min_completion","<","Actor","model","max","sequence","length","-","user_input","+","min_completion","<","Critic","model","max","sequence","length","-","user_input","+","min_completion","<","Reward","model","max","sequence","length","Args",":","config","(","Config",")",":","config","object"]
code: def clean_dataset(config: ConfigType):
      """Clean the datasets by removing too long examples
      The Reward Dataset constraints are:
      - user_input + completion < Reward model max sequence length
      The Actor Dataset constraints are:
      - user_input + completion < Actor model max sequence length
      The RLHF Training Dataset constraints are:
      - user_input + min_completion < Actor model max sequence length
      - user_input + min_completion < Critic model max sequence length
      - user_input + min_completion < Reward model max sequence length
  
      Args:
          config (Config): config object
      """
  
      if isinstance(config, Config):
          print("Start cleaning the dataset for RLHF")
          # constraints
          r_model_max_seq_len = config.reward.max_sequence_length
          a_model_max_seq_len = config.actor.max_sequence_length
          c_model_max_seq_len = config.critic.max_sequence_length
          min_completion = config.actor.min_tokens
          # dataset
          dataset_path = config.trainer.examples_path
          # tokenizers
          r_tokenizer = RewardModel.load_tokenizer(config.reward)
          a_tokenizer = ActorModel.load_tokenizer(config.actor)
          c_tokenizer = CriticModel.load_tokenizer(config.critic)
          # safety tokens
          safety_tokens = config.actor.additonal_prompt_tokens
  
      elif isinstance(config, ConfigActor):
          print("Start cleaning the dataset for Actor")
          # constraint
          a_model_max_seq_len = config.max_sequence_length
          # dataset
          dataset_path = config.train_dataset_path
          # tokenizer
          a_tokenizer = ActorModel.load_tokenizer(config)
          # safety tokens
          safety_tokens = config.additonal_prompt_tokens
  
      elif isinstance(config, ConfigReward):
          print("Start cleaning the dataset for Reward")
          # constraint
          r_model_max_seq_len = config.max_sequence_length
          # dataset
          dataset_path = config.train_dataset_path
          # tokenizer
          r_tokenizer = RewardModel.load_tokenizer(config)
  
      # if there is the datasets
      if os.path.exists(dataset_path):
          # load the dataset
          with open(dataset_path, "r") as f:
              conversations = json.load(f)
  
          # sort in desceding order - longest first
          if isinstance(config, Config):
              conversations = BaseDataset.sort_conversation(
                  conversations,
                  only_input=True,
                  reverse=True,
              )
          else:
              conversations = BaseDataset.sort_conversation(
                  conversations,
                  only_input=False,
                  reverse=True,
              )
  
          old_len = len(conversations)
          # remove too long examples
          # since datasets are ordered by the length
          # we can remove the first elements until we find
          # an example that is not too long
          while len(conversations) > 0:
              # get the text to be tokenized
              if isinstance(config, Config):
                  text = conversations[0]["user_input"]
              else:
                  text = (
                      conversations[0]["user_input"]
                      + conversations[0]["completion"]
                  )
  
              # remove elements from RLHF dataset
              if isinstance(config, Config):
                  a_tokens = a_tokenizer.encode(text, truncation=False)
                  r_tokens = r_tokenizer.encode(text, truncation=False)
                  c_tokens = c_tokenizer.encode(text, truncation=False)
                  if (
                      len(a_tokens) + min_completion + safety_tokens
                      > a_model_max_seq_len
                  ):
                      conversations.pop(0)
                  elif (
                      len(r_tokens) + min_completion + safety_tokens
                      > r_model_max_seq_len
                  ):
                      conversations.pop(0)
                  elif (
                      len(c_tokens) + min_completion + safety_tokens
                      > c_model_max_seq_len
                  ):
                      conversations.pop(0)
                  else:
                      break
  
              # remove elements from Actor dataset
              elif isinstance(config, ConfigActor):
                  tokens = a_tokenizer.encode(text, truncation=False)
                  if len(tokens) + safety_tokens > a_model_max_seq_len:
                      conversations.pop(0)
                  else:
                      break
  
              # remove elements from Reward dataset
              elif isinstance(config, ConfigReward):
                  tokens = r_tokenizer.encode(text, truncation=False)
                  if len(tokens) > r_model_max_seq_len:
                      conversations.pop(0)
                  else:
                      break
  
          # if the number of examples has changed
          if len(conversations) != old_len:
              print("Number of examples before cleaning: ", old_len)
              print(
                  "Number of examples after cleaning: ",
                  len(conversations),
              )
  
              # remove the old dataset
              os.remove(dataset_path)
  
              # save the new dataset
              with open(dataset_path, "w") as f:
                  json.dump(conversations, f, indent=4)
          else:
              print("Dataset is already clean")
  
      else:
          print(
              f"Dataset not found at {dataset_path}"
              f" Skipping cleaning of the dataset"
          )
  
imports: import json
  import os
  import numpy
  from beartype.typing import Dict, List, Union
  from datasets import load_dataset
  from chatllama.rlhf.config import Config, ConfigActor, ConfigReward
  from chatllama.rlhf.reward import RewardModel, CriticModel
  from chatllama.rlhf.actor import ActorModel
label_desc: Use image node_id 4 for calling the BaseDataset obj's underlying member method code with example usage: obj.clean_dataset(config) without return types