Method __init__

utils/dataset.py:164–171

(self, te_dataset, flattened_captions)

class TextEmbeddingDataset:
    def __init__(self, te_dataset, flattened_captions):
        self.te_dataset = te_dataset
        self.flattened_captions = flattened_captions
        self.image_spec_to_te_idx = defaultdict(list)
        # TODO: maybe make this use Dataset object like the latents. But, you won't be caching text embeddings
        # when training on very large datasets, so perhaps it doesn't really matter.
        for i, image_spec in enumerate(flattened_captions['image_spec']):
            self.image_spec_to_te_idx[tuple(image_spec)].append(i)

    def get_text_embeddings(self, image_spec, caption_number):
        return self.te_dataset[self.image_spec_to_te_idx[image_spec][caption_number]]