crfm-helm 0.4.0__py3-none-any.whl → 0.5.10__py3-none-any.whl

helm/clients/image_generation/dalle_mini/data.py

@@ -0,0 +1,442 @@
+import random
+from dataclasses import dataclass, field
+from functools import partial
+from pathlib import Path
+
+import numpy as np
+from datasets import Dataset, load_dataset
+
+from helm.clients.image_generation.dalle_mini.model.text import TextNormalizer
+from helm.common.optional_dependencies import handle_module_not_found_error
+
+try:
+    import jax
+    import jax.numpy as jnp
+    from braceexpand import braceexpand
+except ModuleNotFoundError as e:
+    handle_module_not_found_error(e, ["heim"])
+
+
+@dataclass
+class Dataset:
+    dataset_repo_or_path: str
+    train_file: str = None
+    validation_file: str = None
+    streaming: bool = True
+    use_auth_token: bool = False
+    text_column: str = "caption"
+    encoding_column: str = "encoding"
+    max_train_samples: int = None
+    max_eval_samples: int = None
+    preprocessing_num_workers: int = None
+    overwrite_cache: bool = False
+    do_train: bool = False
+    do_eval: bool = True
+    seed_dataset: int = None
+    shard_by_host: bool = False
+    blank_caption_prob: float = 0.0
+    clip_score_column: str = "clip_score"
+    min_clip_score: float = None
+    max_clip_score: float = None
+    filter_column: str = None
+    filter_value: str = None
+    multi_eval_ds: bool = False
+    train_dataset: Dataset = field(init=False)
+    eval_dataset: Dataset = field(init=False)
+    other_eval_datasets: list = field(init=False)
+    rng_dataset: jnp.ndarray = field(init=False)
+    multi_hosts: bool = field(init=False)
+
+    def __post_init__(self):
+        if self.seed_dataset is None:
+            # create a random seed
+            self.seed_dataset = random.randint(0, 2**32 - 1)
+        # set numpy rng
+        self.np_rng = np.random.default_rng(self.seed_dataset)
+        self.multi_hosts = jax.process_count() > 1
+        # feed blank captions only in streaming mode for now
+        # otherwise dataset could be cached with same blanked captions
+        if self.blank_caption_prob:
+            assert self.streaming is True, "blank_caption_prob can only be used in streaming mode"
+        # define data_files
+        if self.train_file is not None or self.validation_file is not None:
+            # accept braceexpand notation
+            for k in ["train_file", "validation_file"]:
+                f = getattr(self, k)
+                if isinstance(f, str):
+                    setattr(self, k, list(braceexpand(f)))
+            # for list of files, split training data shards by host
+            if isinstance(self.train_file, list) and self.multi_hosts and self.shard_by_host:
+                self.train_file = self.train_file[jax.process_index() :: jax.process_count()]
+            data_files = {
+                "train": self.train_file,
+                "validation": self.validation_file,
+            }
+        else:
+            data_files = None
+
+        # multiple validation datasets
+        if self.multi_eval_ds:
+            assert Path(
+                self.dataset_repo_or_path
+            ).is_dir(), f"{self.dataset_repo_or_path} is not a directory, required for multi_eval_ds"
+            data_files = {
+                split.name: [str(f) for f in split.glob("*.parquet")]
+                for split in Path(self.dataset_repo_or_path).glob("*")
+            }
+            # rename "valid" to "validation" if present for consistency
+            if "valid" in data_files:
+                data_files["validation"] = data_files["valid"]
+                del data_files["valid"]
+            self.dataset_repo_or_path = "parquet"
+
+        # load dataset
+        dataset = load_dataset(
+            self.dataset_repo_or_path,
+            data_files=data_files,
+            streaming=self.streaming,
+            use_auth_token=self.use_auth_token,
+        )
+        if self.do_train:
+            if "train" not in dataset:
+                raise ValueError("Training requires a training dataset")
+            self.train_dataset = dataset["train"]
+            if self.max_train_samples is not None:
+                self.train_dataset = (
+                    self.train_dataset.take(self.max_train_samples)
+                    if self.streaming
+                    else self.train_dataset.select(range(self.max_train_samples))
+                )
+        if self.do_eval:
+            if "validation" not in dataset:
+                raise ValueError("Evaluating requires a validation dataset")
+            self.eval_dataset = dataset["validation"]
+            if self.max_eval_samples is not None:
+                self.eval_dataset = (
+                    self.eval_dataset.take(self.max_eval_samples)
+                    if self.streaming
+                    else self.eval_dataset.select(range(self.max_eval_samples))
+                )
+            # other eval datasets
+            other_eval_splits = dataset.keys() - {"train", "validation"}
+            self.other_eval_datasets = {split: dataset[split] for split in other_eval_splits}
+
+    def preprocess(self, tokenizer, config):
+        # get required config variables
+        decoder_start_token_id = config.decoder_start_token_id
+        normalize_text = config.normalize_text
+        max_length = config.max_text_length
+
+        if self.streaming:
+            # we need to shuffle early in streaming mode
+            if hasattr(self, "train_dataset"):
+                self.train_dataset = self.train_dataset.shuffle(buffer_size=5000, seed=self.seed_dataset)
+        else:
+            self.rng_dataset = jax.random.PRNGKey(self.seed_dataset)
+
+        # filter data
+        partial_filter_function = partial(
+            filter_function,
+            filter_column=self.filter_column,
+            filter_value=self.filter_value,
+            clip_score_column=self.clip_score_column,
+            min_clip_score=self.min_clip_score,
+            max_clip_score=self.max_clip_score,
+        )
+        for ds in ["train_dataset", "eval_dataset"]:
+            if hasattr(self, ds):
+                setattr(
+                    self,
+                    ds,
+                    (
+                        getattr(self, ds).filter(partial_filter_function)
+                        if self.streaming
+                        else getattr(self, ds).filter(
+                            partial_filter_function,
+                            num_proc=self.preprocessing_num_workers,
+                            load_from_cache_file=not self.overwrite_cache,
+                            desc="Filtering datasets",
+                        )
+                    ),
+                )
+        if hasattr(self, "other_eval_datasets"):
+            self.other_eval_datasets = {
+                split: (
+                    ds.filter(partial_filter_function)
+                    if self.streaming
+                    else ds.filter(
+                        partial_filter_function,
+                        num_proc=self.preprocessing_num_workers,
+                        load_from_cache_file=not self.overwrite_cache,
+                        desc="Filtering datasets",
+                    )
+                )
+                for split, ds in self.other_eval_datasets.items()
+            }
+
+        # normalize text
+        if normalize_text:
+            text_normalizer = TextNormalizer()
+            partial_normalize_function = partial(
+                normalize_function,
+                text_column=self.text_column,
+                text_normalizer=text_normalizer,
+            )
+            for ds in ["train_dataset", "eval_dataset"]:
+                if hasattr(self, ds):
+                    setattr(
+                        self,
+                        ds,
+                        (
+                            getattr(self, ds).map(partial_normalize_function)
+                            if self.streaming
+                            else getattr(self, ds).map(
+                                partial_normalize_function,
+                                num_proc=self.preprocessing_num_workers,
+                                load_from_cache_file=not self.overwrite_cache,
+                                desc="Normalizing datasets",
+                            )
+                        ),
+                    )
+            if hasattr(self, "other_eval_datasets"):
+                self.other_eval_datasets = {
+                    split: (
+                        ds.map(partial_normalize_function)
+                        if self.streaming
+                        else ds.map(
+                            partial_normalize_function,
+                            num_proc=self.preprocessing_num_workers,
+                            load_from_cache_file=not self.overwrite_cache,
+                            desc="Normalizing datasets",
+                        )
+                    )
+                    for split, ds in self.other_eval_datasets.items()
+                }
+
+        # blank captions
+        if self.blank_caption_prob:
+            partial_blank_caption_function = partial(
+                blank_caption_function,
+                text_column=self.text_column,
+                blank_caption_prob=self.blank_caption_prob,
+                rng=self.np_rng,
+            )
+            if hasattr(self, "train_dataset"):
+                self.train_dataset = (
+                    self.train_dataset.map(partial_blank_caption_function)
+                    if self.streaming
+                    else self.train_dataset.map(
+                        partial_blank_caption_function,
+                        num_proc=None if self.seed_dataset else self.preprocessing_num_workers,
+                        load_from_cache_file=False,
+                        desc="Blanking some captions",
+                    )
+                )
+
+        # preprocess
+        partial_preprocess_function = partial(
+            preprocess_function,
+            tokenizer=tokenizer,
+            text_column=self.text_column,
+            encoding_column=self.encoding_column,
+            max_length=max_length,
+            decoder_start_token_id=decoder_start_token_id,
+        )
+        for ds in ["train_dataset", "eval_dataset"]:
+            if hasattr(self, ds):
+                setattr(
+                    self,
+                    ds,
+                    (
+                        getattr(self, ds).map(
+                            partial_preprocess_function,
+                            batched=True,
+                            remove_columns=[
+                                self.text_column,
+                                self.encoding_column,
+                            ],
+                        )
+                        if self.streaming
+                        else getattr(self, ds).map(
+                            partial_preprocess_function,
+                            batched=True,
+                            remove_columns=getattr(ds, "column_names"),
+                            num_proc=self.preprocessing_num_workers,
+                            load_from_cache_file=not self.overwrite_cache,
+                            desc="Preprocessing datasets",
+                        )
+                    ),
+                )
+        if hasattr(self, "other_eval_datasets"):
+            self.other_eval_datasets = {
+                split: (
+                    ds.map(
+                        partial_preprocess_function,
+                        batched=True,
+                        remove_columns=[
+                            self.text_column,
+                            self.encoding_column,
+                        ],
+                    )
+                    if self.streaming
+                    else ds.map(
+                        partial_preprocess_function,
+                        batched=True,
+                        remove_columns=getattr(ds, "column_names"),
+                        num_proc=self.preprocessing_num_workers,
+                        load_from_cache_file=not self.overwrite_cache,
+                        desc="Preprocessing datasets",
+                    )
+                )
+                for split, ds in self.other_eval_datasets.items()
+            }
+
+    def dataloader(self, split, batch_size, epoch=None):
+        def _dataloader_datasets_non_streaming(
+            dataset: Dataset,
+            rng: jax.random.PRNGKey = None,
+        ):
+            """
+            Returns batches of size `batch_size` from truncated `dataset`, sharded over all local devices.
+            Shuffle batches if rng is set.
+            """
+            steps_per_epoch = len(dataset) // batch_size
+
+            if rng is not None:
+                batch_idx = jax.random.permutation(rng, len(dataset))
+            else:
+                batch_idx = jnp.arange(len(dataset))
+
+            batch_idx = batch_idx[: steps_per_epoch * batch_size]  # Skip incomplete batch.
+            batch_idx = batch_idx.reshape((steps_per_epoch, batch_size))
+
+            for idx in batch_idx:
+                batch = dataset[idx]
+                batch = {k: jnp.array(v) for k, v in batch.items()}
+                yield batch
+
+        def _dataloader_datasets_streaming(
+            dataset: Dataset,
+            epoch: int,
+        ):
+            keys = ["input_ids", "attention_mask", "labels", "decoder_input_ids"]
+            batch = {k: [] for k in keys}
+            first_loop = True  # stop after one loop in some cases
+            while (self.multi_hosts and split == "train") or first_loop:
+                # in multi-host, we run forever (no epoch) as hosts need to stop
+                # at the same time and training data may not be split equally
+                # For validation data we put the entire batch on each host and then
+                # keep only the one specific to each host (could be improved but not necessary)
+                if epoch is not None:
+                    assert split == "train"
+                    # reshuffle training data at each epoch
+                    dataset.set_epoch(epoch)
+                    epoch += 1
+                for item in dataset:
+                    for k in keys:
+                        batch[k].append(item[k])
+                    if len(batch[keys[0]]) == batch_size:
+                        batch = {k: jnp.array(v) for k, v in batch.items()}
+                        yield batch
+                        batch = {k: [] for k in keys}
+                first_loop = False
+
+        if split == "train":
+            ds = self.train_dataset
+        elif split == "eval":
+            ds = self.eval_dataset
+        else:
+            ds = self.other_eval_datasets[split]
+
+        if self.streaming:
+            return _dataloader_datasets_streaming(ds, epoch)
+        else:
+            if split == "train":
+                self.rng_dataset, input_rng = jax.random.split(self.rng_dataset)
+            return _dataloader_datasets_non_streaming(ds, input_rng)
+
+    @property
+    def length(self):
+        len_train_dataset, len_eval_dataset = None, None
+        if self.streaming:
+            # we don't know the length, let's just assume max_samples if defined
+            if self.max_train_samples is not None:
+                len_train_dataset = self.max_train_samples
+            if self.max_eval_samples is not None:
+                len_eval_dataset = self.max_eval_samples
+        else:
+            len_train_dataset = len(self.train_dataset) if hasattr(self, "train_dataset") else None
+            len_eval_dataset = len(self.eval_dataset) if hasattr(self, "eval_dataset") else None
+        return len_train_dataset, len_eval_dataset
+
+
+def shift_tokens_right(input_ids: np.array, decoder_start_token_id: int):
+    """
+    Shift input ids one token to the right.
+    """
+    shifted_input_ids = np.zeros(input_ids.shape)
+    shifted_input_ids[:, 1:] = input_ids[:, :-1]
+    shifted_input_ids[:, 0] = decoder_start_token_id
+    return shifted_input_ids
+
+
+def blank_caption_function(example, text_column, blank_caption_prob, rng=None):
+    if blank_caption_prob and (rng.random() if rng is not None else np.random.random()) < blank_caption_prob:
+        example[text_column] = ""
+    return example
+
+
+def normalize_function(example, text_column, text_normalizer):
+    example[text_column] = text_normalizer(example[text_column])
+    return example
+
+
+def filter_function(
+    example,
+    min_clip_score,
+    max_clip_score,
+    clip_score_column,
+    filter_column,
+    filter_value,
+):
+    if min_clip_score is not None and example[clip_score_column] < min_clip_score:
+        return False
+    if max_clip_score is not None and example[clip_score_column] > max_clip_score:
+        return False
+    if filter_column is not None and example[filter_column] != filter_value:
+        return False
+    return True
+
+
+def preprocess_function(
+    examples,
+    tokenizer,
+    text_column,
+    encoding_column,
+    max_length,
+    decoder_start_token_id,
+):
+    inputs = examples[text_column]
+    # Setting padding="max_length" as we need fixed length inputs for jitted functions
+    model_inputs = tokenizer(
+        inputs,
+        max_length=max_length,
+        padding="max_length",
+        truncation=True,
+        return_tensors="np",
+    )
+
+    # set up targets
+    # Note: labels correspond to our target indices
+    # decoder input ids are the same but shifted to the right with bos at the beginning (and without last token)
+    labels = examples[encoding_column]
+    labels = np.asarray(labels)
+
+    # We need the labels, in addition to the decoder_input_ids, for the compute_loss function
+    model_inputs["labels"] = labels
+
+    # In our case, this prepends the bos token and removes the last one
+    decoder_input_ids = shift_tokens_right(labels, decoder_start_token_id)
+    model_inputs["decoder_input_ids"] = decoder_input_ids
+
+    return model_inputs

helm/clients/image_generation/dalle_mini/model/__init__.py

@@ -0,0 +1,5 @@
+from helm.clients.image_generation.dalle_mini.model.configuration import DalleBartConfig
+from helm.clients.image_generation.dalle_mini.model.modeling import DalleBart
+from helm.clients.image_generation.dalle_mini.model.partitions import set_partitions
+from helm.clients.image_generation.dalle_mini.model.processor import DalleBartProcessor
+from helm.clients.image_generation.dalle_mini.model.tokenizer import DalleBartTokenizer

helm/clients/image_generation/dalle_mini/model/configuration.py

@@ -0,0 +1,175 @@
+# coding=utf-8
+# Copyright 2021 The Fairseq Authors and The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""DalleBart model configuration"""
+import warnings
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+
+from helm.clients.image_generation.dalle_mini.model.utils import PretrainedFromWandbMixin
+
+logger = logging.get_logger(__name__)
+
+
+class DalleBartConfig(PretrainedFromWandbMixin, PretrainedConfig):
+    model_type = "dallebart"
+    keys_to_ignore_at_inference = ["past_key_values"]
+    attribute_map = {
+        "num_attention_heads": "encoder_attention_heads",
+        "hidden_size": "d_model",
+    }
+
+    def __init__(
+        self,
+        normalize_text=False,
+        encoder_vocab_size=50264,
+        image_vocab_size=16384,  # encoded image token space
+        image_length=256,  # number of encoded tokens
+        max_text_length=64,  # max number of text tokens
+        encoder_layers=12,
+        encoder_ffn_dim=4096,
+        encoder_attention_heads=16,
+        decoder_layers=12,
+        decoder_ffn_dim=4096,
+        decoder_attention_heads=16,
+        activation_function="gelu",
+        d_model=1024,
+        dropout=0.1,
+        attention_dropout=0.0,
+        activation_dropout=0.0,
+        init_std=0.02,
+        scale_embedding=False,
+        gradient_checkpointing=True,
+        use_scan=None,
+        use_cache=True,
+        is_encoder_decoder=True,
+        forced_eos_token_id=None,
+        tie_word_embeddings=False,  # different modalities and sizes
+        do_sample=True,
+        # transformer variants
+        use_bias=False,  # use bias in attention and dense layers (except for lm_head)
+        ln_type="layernorm",  # layer normalization type, "rmsnorm", "layernorm"
+        ln_positions="normformer",  # layer normalization positions, "normformer", "swinv2", "cogview", "postln", "preln", "deepnet" (same as postln), "subln"
+        use_head_scale=False,  # used in NormFormer
+        use_cosine_attention=False,  # used in Swin v2
+        tau_init=0.05,  # used only in cosine attention (Swin v2)
+        use_absolute_position_embeddings=True,  # default
+        use_swin_position_embeddings=False,  # used in Swin v1/v2
+        use_deepnet_scaling=False,  # used in Deepnet
+        use_subln_init=False,
+        use_glu=True,  # "GLU Variants Improve Transformer"
+        use_alibi=False,  # Not implemented yet - from "Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation"
+        sinkhorn_iters=1,  # used in SinkFormers
+        use_final_ln_encoder=True,  # final layer normalization in encoder
+        use_final_ln_decoder=True,  # final layer normalization in decoder
+        # parameters that should not be necessary but could affect results
+        force_ln_scale=False,  # force scale in layernorm even when followed by dense layers
+        **kwargs,
+    ):
+        # text normalizer
+        self.normalize_text = normalize_text
+
+        # transformer variants
+        self.use_bias = use_bias
+        assert ln_type in [
+            "rmsnorm",
+            "layernorm",
+        ], "ln_type must be 'rmsnorm' or 'layernorm'"
+        self.ln_type = ln_type
+        if ln_positions == "deepnet":
+            ln_positions = "postln"
+        assert ln_positions in [
+            "normformer",
+            "swinv2",
+            "cogview",
+            "postln",
+            "preln",
+            "subln",
+        ], "ln_positions must be 'normformer', 'swinv2', 'cogview', 'postln', 'preln', 'subln'"
+        self.use_head_scale = use_head_scale
+        assert use_alibi is False, "use_alibi is not supported yet"
+        self.ln_positions = ln_positions
+        self.use_cosine_attention = use_cosine_attention
+        self.tau_init = tau_init
+        self.use_absolute_position_embeddings = use_absolute_position_embeddings
+        self.use_swin_position_embeddings = use_swin_position_embeddings
+        self.use_deepnet_scaling = use_deepnet_scaling
+        self.use_subln_init = use_subln_init
+        self.use_glu = use_glu
+        self.use_alibi = use_alibi
+        self.sinkhorn_iters = sinkhorn_iters
+        if ln_positions == "postln":
+            assert use_final_ln_encoder, "use_final_ln_encoder must be True when ln_positions is 'postln'"
+            assert use_final_ln_decoder, "use_final_ln_decoder must be True when ln_positions is 'postln'"
+        self.use_final_ln_encoder = use_final_ln_encoder
+        self.use_final_ln_decoder = use_final_ln_decoder
+        self.force_ln_scale = force_ln_scale
+
+        # common parameters
+        self.encoder_vocab_size = encoder_vocab_size
+        self.image_vocab_size = image_vocab_size
+        self.image_length = image_length
+        self.max_text_length = max_text_length
+        self.d_model = d_model
+        self.encoder_ffn_dim = encoder_ffn_dim
+        self.encoder_layers = encoder_layers
+        self.encoder_attention_heads = encoder_attention_heads
+        self.decoder_ffn_dim = decoder_ffn_dim
+        self.decoder_layers = decoder_layers
+        self.decoder_attention_heads = decoder_attention_heads
+        self.dropout = dropout
+        self.attention_dropout = attention_dropout
+        self.activation_dropout = activation_dropout
+        self.activation_function = activation_function
+        self.init_std = init_std
+        self.use_cache = use_cache
+        self.gradient_checkpointing = gradient_checkpointing
+        # all layers are the same in most configurations
+        self.use_scan = use_scan if use_scan is not None else ln_positions != "swinv2"
+        assert not (self.use_scan and ln_positions == "swinv2"), "scan cannot be used with 'swinv2'"
+        self.scale_embedding = scale_embedding  # scale factor will be sqrt(d_model) if True
+
+        # special token id's are appended to vocab if not provided
+        decoder_start_token_id = kwargs.pop("decoder_start_token_id", image_vocab_size)
+        bos_token_id = kwargs.pop("bos_token_id", image_vocab_size)
+        pad_token_id = kwargs.pop("pad_token_id", image_vocab_size)
+        eos_token_id = kwargs.pop("eos_token_id", image_vocab_size)
+
+        # we generate to image_length + 1 (for bos) by default
+        min_length = kwargs.pop("min_length", image_length + 1)
+        max_length = kwargs.pop("max_length", image_length + 1)
+
+        super().__init__(
+            # args required in parent class
+            is_encoder_decoder=is_encoder_decoder,
+            tie_word_embeddings=tie_word_embeddings,
+            forced_eos_token_id=forced_eos_token_id,
+            decoder_start_token_id=decoder_start_token_id,
+            bos_token_id=bos_token_id,
+            pad_token_id=pad_token_id,
+            eos_token_id=eos_token_id,
+            min_length=min_length,
+            max_length=max_length,
+            do_sample=do_sample,
+            **kwargs,
+        )
+
+        # ensure backward compatibility for BART CNN models
+        if self.forced_bos_token_id is None and kwargs.get("force_bos_token_to_be_generated", False):
+            self.forced_bos_token_id = self.bos_token_id
+            warnings.warn(
+                f"Please make sure the config includes `forced_bos_token_id={self.bos_token_id}` in future versions."
+                "The config can simply be saved and uploaded again to be fixed."
+            )