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
class_docstr_tok: ["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["prompt"] + 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["prompt"] + 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 RmStaticDataset object from inherited base classes: Dataset with example: obj = RmStaticDataset(dataset, tokenizer, max_length, special_token)