How to include sample weight when building dataset

Hi,
I try to build custom dataset that contain three element (image, label, sample_weight). I was able to generate only 2 elements, but fail when adding the third element. sample_weight value just simply derived from class label

class CustomDataset(tfds.core.GeneratorBasedBuilder):

  VERSION = tfds.core.Version('1.0.8')
  RELEASE_NOTES = {
      '1.0.8': 'add sample_weight',
  }
  
  LBL = dict(zip(['A', 'B', 'C', 'D', 'E',
                  'F', 'G', 'H', 'I', 'K',
                  'L', 'M', 'N', 'O', 'P',
                  'Q', 'R', 'S', 'T',
                  'U', 'V', 'W', 'X', 'Y',
                  'Z', 'AA', 'BB', 'CC'], range(28)))
    
  weights = dict(zip(['A', 'B', 'C', 'D', 'E',
                  'F', 'G', 'H', 'I', 'K',
                  'L', 'M', 'N', 'O', 'P',
                  'Q', 'R', 'S', 'T',
                  'U', 'V', 'W', 'X', 'Y',
                  'Z', 'AA', 'BB', 'CC'], [1,1,1,1,1,1,1,8.94,1,1,1,1,1,8.94,1,1,1,1,1,1,1,1,1,1,8.94,1,1,1]))
  
  def _info(self) -> tfds.core.DatasetInfo:
    """Returns the dataset metadata."""
    return tfds.core.DatasetInfo(
        builder=self,
        description=_DESCRIPTION,
        features=tfds.features.FeaturesDict({

            'image': tfds.features.Image(shape=(224, 224, 3)),
            'label': tfds.features.ClassLabel(num_classes=28),
        }),

        supervised_keys=('image', 'label'),
        citation="",
    )

  def _split_generators(self, dl_manager: tfds.download.DownloadManager):
    """Returns SplitGenerators."""

    archive_path = dl_manager.manual_dir / 'train.zip'
    extracted_path_train = dl_manager.extract(archive_path)

    archive_path_val = dl_manager.manual_dir / 'val.zip'
    extracted_path_val = dl_manager.extract(archive_path_val)

    return {
        'train': self._generate_examples(extracted_path_train),
        'val': self._generate_examples(extracted_path_val),
    }

  def _generate_examples(self, path):
    """Yields examples."""
    for f in path.glob('**/*.png'):
         key = basename(f)
         lbl = basename(os.path.split(f)[0])
         yield key, {

                  'image': f,
                  'label': self.LBL[lbl],
                  'sample_weight': self.weights[lbl]
         }

I tried adding sample_weight in _info but error was raised due to assert len(supervised_keys) == 2 in tfds.core.dataset_info

Start with the custom builder's _info method and the supervised_keys assertion in tfds.core.dataset_info. Review how _generate_examples yields image, label, and sample_weight, then determine the supported metadata shape and expected behavior for supervised data with three elements. Done means the builder can represent and load sample weights without triggering the assertion.