class_docstr: Dataset for reward model
  
  Args:
      dataset: dataset for reward model
      tokenizer: tokenizer for reward model
      max_length: max length of input
      special_token: special token at the end of sentence
code: def __init__(
      self,
      dataset,
      tokenizer: Callable,
      max_length: int,
      special_token=None,
  ) -> None:
      super().__init__()
      self.end_token = (
          tokenizer.eos_token
          if special_token is None
          else special_token
      )
  
      chosen = [
          data["chosen"] + self.end_token
          for data in tqdm(dataset, disable=not is_rank_0())
      ]
      chosen_token = tokenizer(
          chosen,
          max_length=max_length,
          padding="max_length",
          truncation=True,
          return_tensors="pt",
      )
      self.chosen = {
          "input_ids": chosen_token["input_ids"],
          "attention_mask": chosen_token["attention_mask"],
      }
  
      reject = [
          data["rejected"] + self.end_token
          for data in tqdm(dataset, disable=not is_rank_0())
      ]
      reject_token = tokenizer(
          reject,
          max_length=max_length,
          padding="max_length",
          truncation=True,
          return_tensors="pt",
      )
      self.reject = {
          "input_ids": reject_token["input_ids"],
          "attention_mask": reject_token["attention_mask"],
      }
  
file_path: colossalai/applications/Chat/coati/dataset/reward_dataset.py
imports: from typing import Callable
  from torch.utils.data import Dataset
  from tqdm import tqdm
  from .utils import is_rank_0
label_desc: Use image node_id 1 to create a new HhRlhfDataset object from inherited base classes: Dataset with example: obj = HhRlhfDataset(dataset, tokenizer, max_length, special_token)