content: from .base_dataset import BaseDataset
  import sys
  import random
  
  
  class NLVR2Dataset(BaseDataset):
      def __init__(self, *args, split="", **kwargs):
          assert split in ["train", "val", "test"]
          self.split = split
  
          if split == "train":
              names = ["nlvr2_train"]
          elif split == "val":
              names = ["nlvr2_dev", "nlvr2_test1"]
          elif split == "test":
              names = ["nlvr2_dev", "nlvr2_test1"]
  
          super().__init__(
              *args,
              **kwargs,
              names=names,
              text_column_name="questions",
              remove_duplicate=False,
          )
  
      def __getitem__(self, index):
          result = None
          while result is None:
              try:
                  image_tensor_0 = self.get_image(index, image_key="image_0")["image"]
                  image_tensor_1 = self.get_image(index, image_key="image_1")["image"]
                  text = self.get_text(index)["text"]
                  result = True
              except:
                  print(
                      f"error while read file idx {index} in {self.names[0]}",
                      file=sys.stderr,
                  )
                  index = random.randint(0, len(self.index_mapper) - 1)
  
          index, question_index = self.index_mapper[index]
          answers = self.table["answers"][index][question_index].as_py()
          answers = answers == "True"
  
          return {
              "image_0": image_tensor_0,
              "image_1": image_tensor_1,
              "text": text,
              "answers": answers,
              "table_name": self.table_names[index],
          }
  