Function _cache_text_embeddings

utils/dataset.py:177–201

(metadata_dataset, map_fn, i, cache_dir, regenerate_cache, caching_batch_size)

def _cache_text_embeddings(metadata_dataset, map_fn, i, cache_dir, regenerate_cache, caching_batch_size):

    def flatten_captions(example):
        result = {key: [] for key in example}
        for i, captions in enumerate(example['caption']):
            for caption in captions:
                result['caption'].append(caption)
                for key, value in example.items():
                    if key == 'caption':
                        continue
                    result[key].append(value[i])
        return result

    flattened_captions = metadata_dataset.map(flatten_captions, batched=True, keep_in_memory=True, remove_columns=metadata_dataset.column_names)
    te_dataset = _map_and_cache(
        flattened_captions,
        map_fn,
        cache_dir,
        cache_file_prefix=f'text_embeddings_{i}_',
        new_fingerprint_args=[i],
        regenerate_cache=regenerate_cache,
        caching_batch_size=caching_batch_size,
    )
    assert len(te_dataset) == len(flattened_captions)
    return TextEmbeddingDataset(te_dataset, flattened_captions)


# The smallest unit of a dataset. Represents a single size bucket from a single folder of images