Method cache_text_embeddings

utils/dataset.py:300–303

(self, map_fn, i, regenerate_cache=False, caching_batch_size=1)

        self.iteration_order = datasets.load_from_disk(str(iteration_order_cache_dir))

    def cache_text_embeddings(self, map_fn, i, regenerate_cache=False, caching_batch_size=1):
        print(f'caching text embeddings: {self.size_bucket}')
        te_dataset = _cache_text_embeddings(self.metadata_dataset, map_fn, i, self.cache_dir, regenerate_cache, caching_batch_size)
        self.text_embedding_datasets.append(te_dataset)

    def add_text_embedding_dataset(self, te_dataset):
        self.text_embedding_datasets.append(te_dataset)