code: def test_dataloader_has_correct_keys(df):
      cfg = ConfigProblemBase(
          dataset=ConfigDPODataset(
              prompt_column=("prompt_column",),
              answer_column="answer_column",
              rejected_answer_column="rejected_answer_column",
              parent_id_column="None",
          )
      )
  
      dataset = CustomDataset(df, cfg, mode="train")
      dataloader = torch.utils.data.DataLoader(
          dataset, batch_size=16, shuffle=True
      )
  
      for idx, batch in tqdm(
          enumerate(dataloader), total=len(dataloader)
      ):
          for key in batch:
              if idx != len(dataloader) - 1:
                  assert batch[key].size(0) == 16, (
                      key,
                      batch[key].shape,
                  )
  
              keys = [
                  "chosen_input_ids",
                  "chosen_attention_mask",
                  "chosen_labels",
                  "rejected_input_ids",
                  "rejected_attention_mask",
                  "rejected_labels",
                  "prompt_input_ids",
                  "prompt_attention_mask",
              ]
              assert set(batch.keys()) - set(keys) == set()
              assert set(keys) - set(batch.keys()) == set()
  
file_path: h2o-llmstudio/tests/src/datasets/test_text_dpo_modeling_ds.py
imports: import numpy
  import pandas
  import pytest
  import torch
  from tqdm import tqdm
  from llm_studio.python_configs.text_causal_language_modeling_config import ConfigNLPCausalLMTokenizer
  from llm_studio.python_configs.text_dpo_modeling_config import ConfigDPODataset, ConfigProblemBase
  from llm_studio.src.datasets.text_dpo_modeling_ds import CustomDataset
label_desc: Use image node_id 5 for calling a global function with example usage: test_dataloader_has_correct_keys(df) without return types