code: def __init__(
      self,
      train_file: str,
      tokenizer: AutoTokenizer,
      special_token=None,
      max_length=512,
  ) -> None:
      super(EasyRewardDataset, self).__init__()
      self.chosen = []
      self.reject = []
      if special_token is None:
          self.end_token = tokenizer.eos_token
      else:
          self.end_token = special_token
      print(self.end_token)
      # read all lines in the train_file to a list
      with open(train_file, "r", encoding="UTF-8") as f:
          all_lines = f.readlines()
      for line in tqdm(all_lines):
          data = json.loads(line)
          prompt = "提问：" + data["prompt"] + " 回答："
  
          chosen = prompt + data["chosen"] + self.end_token
          chosen_token = tokenizer(
              chosen,
              max_length=max_length,
              padding="max_length",
              truncation=True,
              return_tensors="pt",
          )
          self.chosen.append(
              {
                  "input_ids": chosen_token["input_ids"],
                  "attention_mask": chosen_token["attention_mask"],
              }
          )
  
          reject = prompt + data["rejected"] + self.end_token
          reject_token = tokenizer(
              reject,
              max_length=max_length,
              padding="max_length",
              truncation=True,
              return_tensors="pt",
          )
          self.reject.append(
              {
                  "input_ids": reject_token["input_ids"],
                  "attention_mask": reject_token["attention_mask"],
              }
          )
  
imports: import copy
  import json
  from typing import Dict, Sequence
  import torch
  from torch.utils.data import Dataset
  from tqdm import tqdm
  from transformers import AutoTokenizer
label_desc: Use image node_id 1 to create a new EasyRewardDataset object from inherited base classes: Dataset with example: obj = EasyRewardDataset(train_file, tokenizer, special_token, max_length)