code: def create_dataset_split(
      current_dataset,
      raw_dataset,
      train_phase,
      tokenizer,
      end_of_conversation_token,
      max_seq_len,
  ):
      prompt_dataset = []
      chosen_dataset = []
      reject_dataset = []
      if train_phase == 1:
          for i, tmp_data in enumerate(current_dataset):
              # tokenize the text
              chosen_sentence = raw_dataset.get_prompt_and_chosen(
                  tmp_data
              )  # the accept response
              if chosen_sentence is not None:
                  chosen_sentence += end_of_conversation_token
                  chosen_token = tokenizer(
                      chosen_sentence,
                      max_length=max_seq_len,
                      padding="max_length",
                      truncation=True,
                      return_tensors="pt",
                  )
                  chosen_token["input_ids"] = chosen_token[
                      "input_ids"
                  ].squeeze(0)
                  chosen_token["attention_mask"] = chosen_token[
                      "attention_mask"
                  ].squeeze(0)
                  chosen_dataset.append(chosen_token)
          print(
              f"Creating dataset {raw_dataset.dataset_name_clean} for {train_phase=} size={len(chosen_dataset)}"
          )
  
      elif train_phase == 2:
          for i, tmp_data in enumerate(current_dataset):
              # tokenize the text
              chosen_sentence = raw_dataset.get_prompt_and_chosen(
                  tmp_data
              )  # the accept response
              reject_sentence = raw_dataset.get_prompt_and_rejected(
                  tmp_data
              )  # the accept response
              if (
                  chosen_sentence is not None
                  and reject_sentence is not None
              ):
                  chosen_sentence += (
                      end_of_conversation_token  # the accept response
                  )
                  reject_sentence += end_of_conversation_token
                  chosen_token = tokenizer(
                      chosen_sentence,
                      max_length=max_seq_len,
                      padding="max_length",
                      truncation=True,
                      return_tensors="pt",
                  )
                  reject_token = tokenizer(
                      reject_sentence,
                      max_length=max_seq_len,
                      padding="max_length",
                      truncation=True,
                      return_tensors="pt",
                  )
                  chosen_token["input_ids"] = chosen_token["input_ids"]
                  chosen_token["attention_mask"] = chosen_token[
                      "attention_mask"
                  ]
                  chosen_dataset.append(chosen_token)
  
                  reject_token["input_ids"] = reject_token["input_ids"]
                  reject_token["attention_mask"] = reject_token[
                      "attention_mask"
                  ]
                  reject_dataset.append(reject_token)
          print(
              f"Creating dataset {raw_dataset.dataset_name_clean} for {train_phase=} size={len(chosen_dataset)}"
          )
  
      elif train_phase == 3:
          filtered = 0
          for i, tmp_data in enumerate(current_dataset):
              # tokenize the text
              prompt = raw_dataset.get_prompt(tmp_data)
              if prompt is not None:
                  prompt_token = tokenizer(prompt, return_tensors="pt")
                  if (
                      prompt_token["input_ids"].size()[-1]
                      <= max_seq_len
                  ):
                      for key_word in ["input_ids", "attention_mask"]:
                          prompt_token[key_word] = (
                              prompt_token[key_word].squeeze(0).flip(0)
                          )
                      prompt_dataset.append(prompt_token)
                  else:
                      filtered += 1
          print(
              f"Creating dataset {raw_dataset.dataset_name_clean} "
              f"for {train_phase=} size={len(prompt_dataset)} {filtered=}"
          )
  
      return PromptDataset(
          prompt_dataset,
          chosen_dataset,
          reject_dataset,
          tokenizer.pad_token_id,
          train_phase,
      )
  
imports: import torch
  from torch.utils.data import Dataset, Subset, ConcatDataset
  from torch.nn.utils.rnn import pad_sequence
  import torch.nn.functional
  from datasets import load_dataset
  import numpy
  import os
  import hashlib
  from itertools import chain
  from dschat.utils.data import raw_datasets
  from deepspeed.accelerator import get_accelerator
label_desc: Use image node_id 4 for calling a global function with example usage: create_dataset_split(current_dataset, raw_dataset, train_phase, tokenizer, end_of_conversation_token, max_seq_len) and returns: PromptDataset