tom-doerr/logo_generator

1import random2from dataclasses import dataclass, field3from functools import partial4 5import jax6import jax.numpy as jnp7import numpy as np8from braceexpand import braceexpand9from datasets import Dataset, load_dataset10 11from .model.text import TextNormalizer12 13 14@dataclass15class Dataset:16    dataset_repo_or_path: str17    train_file: str = None18    validation_file: str = None19    streaming: bool = True20    use_auth_token: bool = False21    text_column: str = "caption"22    encoding_column: str = "encoding"23    max_train_samples: int = None24    max_eval_samples: int = None25    preprocessing_num_workers: int = None26    overwrite_cache: bool = False27    do_train: bool = False28    do_eval: bool = True29    seed_dataset: int = None30    shard_by_host: bool = False31    blank_caption_prob: float = 0.032    clip_score_column: str = "clip_score"33    min_clip_score: float = None34    max_clip_score: float = None35    filter_column: str = None36    filter_value: str = None37    train_dataset: Dataset = field(init=False)38    eval_dataset: Dataset = field(init=False)39    rng_dataset: jnp.ndarray = field(init=False)40    multi_hosts: bool = field(init=False)41 42    def __post_init__(self):43        if self.seed_dataset is None:44            # create a random seed45            self.seed_dataset = random.randint(0, 2**32 - 1)46        # set numpy rng47        self.np_rng = np.random.default_rng(self.seed_dataset)48        self.multi_hosts = jax.process_count() > 149        # feed blank captions only in streaming mode for now50        # otherwise dataset could be cached with same blanked captions51        if self.blank_caption_prob:52            assert (53                self.streaming is True54            ), "blank_caption_prob can only be used in streaming mode"55        # define data_files56        if self.train_file is not None or self.validation_file is not None:57            # accept braceexpand notation58            for k in ["train_file", "validation_file"]:59                f = getattr(self, k)60                if isinstance(f, str):61                    setattr(self, k, list(braceexpand(f)))62            # for list of files, split training data shards by host63            if (64                isinstance(self.train_file, list)65                and self.multi_hosts66                and self.shard_by_host67            ):68                self.train_file = self.train_file[69                    jax.process_index() :: jax.process_count()70                ]71            data_files = {72                "train": self.train_file,73                "validation": self.validation_file,74            }75        else:76            data_files = None77 78        # load dataset79        dataset = load_dataset(80            self.dataset_repo_or_path,81            data_files=data_files,82            streaming=self.streaming,83            use_auth_token=self.use_auth_token,84        )85        if self.do_train:86            if "train" not in dataset:87                raise ValueError("Training requires a training dataset")88            self.train_dataset = dataset["train"]89            if self.max_train_samples is not None:90                self.train_dataset = (91                    self.train_dataset.take(self.max_train_samples)92                    if self.streaming93                    else self.train_dataset.select(range(self.max_train_samples))94                )95        if self.do_eval:96            if "validation" not in dataset:97                raise ValueError("Evaluating requires a validation dataset")98            self.eval_dataset = dataset["validation"]99            if self.max_eval_samples is not None:100                self.eval_dataset = (101                    self.eval_dataset.take(self.max_eval_samples)102                    if self.streaming103                    else self.eval_dataset.select(range(self.max_eval_samples))104                )105 106    def preprocess(self, tokenizer, config):107        # get required config variables108        decoder_start_token_id = config.decoder_start_token_id109        normalize_text = config.normalize_text110        max_length = config.max_text_length111 112        if self.streaming:113            # we need to shuffle early in streaming mode114            if hasattr(self, "train_dataset"):115                self.train_dataset = self.train_dataset.shuffle(116                    buffer_size=5000, seed=self.seed_dataset117                )118        else:119            self.rng_dataset = jax.random.PRNGKey(self.seed_dataset)120 121        # filter data122        partial_filter_function = partial(123            filter_function,124            filter_column=self.filter_column,125            filter_value=self.filter_value,126            clip_score_column=self.clip_score_column,127            min_clip_score=self.min_clip_score,128            max_clip_score=self.max_clip_score,129        )130        for ds in ["train_dataset", "eval_dataset"]:131            if hasattr(self, ds):132                setattr(133                    self,134                    ds,135                    (136                        getattr(self, ds).filter(partial_filter_function)137                        if self.streaming138                        else getattr(self, ds).filter(139                            partial_filter_function,140                            num_proc=self.preprocessing_num_workers,141                            load_from_cache_file=not self.overwrite_cache,142                            desc="Filtering datasets",143                        )144                    ),145                )146 147        # normalize text148        if normalize_text:149            text_normalizer = TextNormalizer()150            partial_normalize_function = partial(151                normalize_function,152                text_column=self.text_column,153                text_normalizer=text_normalizer,154            )155            for ds in ["train_dataset", "eval_dataset"]:156                if hasattr(self, ds):157                    setattr(158                        self,159                        ds,160                        (161                            getattr(self, ds).map(partial_normalize_function)162                            if self.streaming163                            else getattr(self, ds).map(164                                partial_normalize_function,165                                num_proc=self.preprocessing_num_workers,166                                load_from_cache_file=not self.overwrite_cache,167                                desc="Normalizing datasets",168                            )169                        ),170                    )171 172        # blank captions173        if self.blank_caption_prob:174            partial_blank_caption_function = partial(175                blank_caption_function,176                text_column=self.text_column,177                blank_caption_prob=self.blank_caption_prob,178                rng=self.np_rng,179            )180            if hasattr(self, "train_dataset"):181                self.train_dataset = (182                    self.train_dataset.map(partial_blank_caption_function)183                    if self.streaming184                    else self.train_dataset.map(185                        partial_blank_caption_function,186                        num_proc=None187                        if self.seed_dataset188                        else self.preprocessing_num_workers,189                        load_from_cache_file=False,190                        desc="Blanking some captions",191                    )192                )193 194        # preprocess195        partial_preprocess_function = partial(196            preprocess_function,197            tokenizer=tokenizer,198            text_column=self.text_column,199            encoding_column=self.encoding_column,200            max_length=max_length,201            decoder_start_token_id=decoder_start_token_id,202        )203        for ds in ["train_dataset", "eval_dataset"]:204            if hasattr(self, ds):205                setattr(206                    self,207                    ds,208                    (209                        getattr(self, ds).map(210                            partial_preprocess_function,211                            batched=True,212                            remove_columns=[213                                self.text_column,214                                self.encoding_column,215                            ],216                        )217                        if self.streaming218                        else getattr(self, ds).map(219                            partial_preprocess_function,220                            batched=True,221                            remove_columns=getattr(ds, "column_names"),222                            num_proc=self.preprocessing_num_workers,223                            load_from_cache_file=not self.overwrite_cache,224                            desc="Preprocessing datasets",225                        )226                    ),227                )228 229    def dataloader(self, split, batch_size, epoch=None):230        def _dataloader_datasets_non_streaming(231            dataset: Dataset,232            rng: jax.random.PRNGKey = None,233        ):234            """235            Returns batches of size `batch_size` from truncated `dataset`, sharded over all local devices.236            Shuffle batches if rng is set.237            """238            steps_per_epoch = len(dataset) // batch_size239 240            if rng is not None:241                batch_idx = jax.random.permutation(rng, len(dataset))242            else:243                batch_idx = jnp.arange(len(dataset))244 245            batch_idx = batch_idx[246                : steps_per_epoch * batch_size247            ]  # Skip incomplete batch.248            batch_idx = batch_idx.reshape((steps_per_epoch, batch_size))249 250            for idx in batch_idx:251                batch = dataset[idx]252                batch = {k: jnp.array(v) for k, v in batch.items()}253                yield batch254 255        def _dataloader_datasets_streaming(256            dataset: Dataset,257            epoch: int,258        ):259            keys = ["input_ids", "attention_mask", "labels", "decoder_input_ids"]260            batch = {k: [] for k in keys}261            first_loop = True  # stop after one loop in some cases262            while (self.multi_hosts and split == "train") or first_loop:263                # in multi-host, we run forever (no epoch) as hosts need to stop264                # at the same time and training data may not be split equally265                # For validation data we put the entire batch on each host and then266                # keep only the one specific to each host (could be improved but not necessary)267                if epoch is not None:268                    assert split == "train"269                    # reshuffle training data at each epoch270                    dataset.set_epoch(epoch)271                    epoch += 1272                for item in dataset:273                    for k in keys:274                        batch[k].append(item[k])275                    if len(batch[keys[0]]) == batch_size:276                        batch = {k: jnp.array(v) for k, v in batch.items()}277                        yield batch278                        batch = {k: [] for k in keys}279                first_loop = False280 281        if split == "train":282            ds = self.train_dataset283        elif split == "eval":284            ds = self.eval_dataset285        else:286            raise ValueError(f'split must be "train" or "eval", got {split}')287 288        if self.streaming:289            return _dataloader_datasets_streaming(ds, epoch)290        else:291            if split == "train":292                self.rng_dataset, input_rng = jax.random.split(self.rng_dataset)293            return _dataloader_datasets_non_streaming(ds, input_rng)294 295    @property296    def length(self):297        len_train_dataset, len_eval_dataset = None, None298        if self.streaming:299            # we don't know the length, let's just assume max_samples if defined300            if self.max_train_samples is not None:301                len_train_dataset = self.max_train_samples302            if self.max_eval_samples is not None:303                len_eval_dataset = self.max_eval_samples304        else:305            len_train_dataset = (306                len(self.train_dataset) if hasattr(self, "train_dataset") else None307            )308            len_eval_dataset = (309                len(self.eval_dataset) if hasattr(self, "eval_dataset") else None310            )311        return len_train_dataset, len_eval_dataset312 313 314def shift_tokens_right(input_ids: np.array, decoder_start_token_id: int):315    """316    Shift input ids one token to the right.317    """318    shifted_input_ids = np.zeros(input_ids.shape)319    shifted_input_ids[:, 1:] = input_ids[:, :-1]320    shifted_input_ids[:, 0] = decoder_start_token_id321    return shifted_input_ids322 323 324def blank_caption_function(example, text_column, blank_caption_prob, rng=None):325    if (326        blank_caption_prob327        and (rng.random() if rng is not None else np.random.random())328        < blank_caption_prob329    ):330        example[text_column] = ""331    return example332 333 334def normalize_function(example, text_column, text_normalizer):335    example[text_column] = text_normalizer(example[text_column])336    return example337 338 339def filter_function(340    example,341    min_clip_score,342    max_clip_score,343    clip_score_column,344    filter_column,345    filter_value,346):347    if min_clip_score is not None and example[clip_score_column] < min_clip_score:348        return False349    if max_clip_score is not None and example[clip_score_column] > max_clip_score:350        return False351    if filter_column is not None and example[filter_column] != filter_value:352        return False353    return True354 355 356def preprocess_function(357    examples,358    tokenizer,359    text_column,360    encoding_column,361    max_length,362    decoder_start_token_id,363):364    inputs = examples[text_column]365    # Setting padding="max_length" as we need fixed length inputs for jitted functions366    model_inputs = tokenizer(367        inputs,368        max_length=max_length,369        padding="max_length",370        truncation=True,371        return_tensors="np",372    )373 374    # set up targets375    # Note: labels correspond to our target indices376    # decoder input ids are the same but shifted to the right with bos at the beginning (and without last token)377    labels = examples[encoding_column]378    labels = np.asarray(labels)379 380    # We need the labels, in addition to the decoder_input_ids, for the compute_loss function381    model_inputs["labels"] = labels382 383    # In our case, this prepends the bos token and removes the last one384    decoder_input_ids = shift_tokens_right(labels, decoder_start_token_id)385    model_inputs["decoder_input_ids"] = decoder_input_ids386 387    return model_inputs388