Format with black and isort and lint with flake8

README.md

@@ -7,18 +7,22 @@ tags:
 - vit
 ---
 # CLIP-Spanish
+
 CLIP Spanish is a CLIP-like model for Spanish language. It is composed of [BERTIN](https://huggingface.co/bertin-project/bertin-roberta-base-spanish) as a language encoder and the ViT-B/32 image encoder from [CLIP](https://huggingface.co/openai/clip-vit-base-patch32). The model is implemented in [Flax](https://github.com/google/flax), including training scripts (see `training.md`).
 This is part of the [Flax/Jax Community Week](https://discuss.huggingface.co/t/open-to-the-community-community-week-using-jax-flax-for-nlp-cv/7104), organised by [HuggingFace](https://huggingface.co/) and TPU usage sponsored by Google.
 
 ## Spanish WIT
+
 We used a subset of 141,230 Spanish captions from the [WIT dataset](https://github.com/google-research-datasets/wit) for training.
 
 ## Team members
+
 - Eduardo González Ponferrada ([edugp](https://huggingface.co/edugp))
 - Manu Romero ([mrm8488](https://huggingface.co/))
 - María Grandury ([mariagrandury](https://huggingface.co/))
 
 ## Useful links
+
 - [Community Week timeline](https://discuss.huggingface.co/t/open-to-the-community-community-week-using-jax-flax-for-nlp-cv/7104#summary-timeline-calendar-6)
 - [Community Week README](https://github.com/huggingface/transformers/blob/master/examples/research_projects/jax-projects/README.md)
 - [Community Week thread](https://discuss.huggingface.co/t/bertin-pretrain-roberta-large-from-scratch-in-spanish/7125)

configuration_hybrid_clip.py

@@ -3,7 +3,6 @@ import copy
 from transformers.configuration_utils import PretrainedConfig
 from transformers.utils import logging
 
-
 logger = logging.get_logger(__name__)
 
 
@@ -64,19 +63,25 @@ class HybridCLIPConfig(PretrainedConfig):
         self.text_config = AutoConfig.for_model(text_model_type, **text_config)
 
         if vision_model_type == "clip":
-            self.vision_config = AutoConfig.for_model(vision_model_type, **vision_config).vision_config
+            self.vision_config = AutoConfig.for_model(
+                vision_model_type, **vision_config
+            ).vision_config
         elif vision_model_type == "clip_vision_model":
             from transformers import CLIPVisionConfig
 
             self.vision_config = CLIPVisionConfig(**vision_config)
         else:
-            self.vision_config = AutoConfig.for_model(vision_model_type, **vision_config)
+            self.vision_config = AutoConfig.for_model(
+                vision_model_type, **vision_config
+            )
 
         self.projection_dim = projection_dim
         self.initializer_factor = 1.0
 
     @classmethod
-    def from_text_vision_configs(cls, text_config: PretrainedConfig, vision_config: PretrainedConfig, **kwargs):
+    def from_text_vision_configs(
+        cls, text_config: PretrainedConfig, vision_config: PretrainedConfig, **kwargs
+    ):
         r"""
         Instantiate a :class:`HybridCLIPConfig` (or a derived class) from text model configuration and
         vision model configuration.
@@ -84,7 +89,11 @@ class HybridCLIPConfig(PretrainedConfig):
             :class:`HybridCLIPConfig`: An instance of a configuration object
         """
 
-        return cls(text_config=text_config.to_dict(), vision_config=vision_config.to_dict(), **kwargs)
+        return cls(
+            text_config=text_config.to_dict(),
+            vision_config=vision_config.to_dict(),
+            **kwargs
+        )
 
     def to_dict(self):
         """

discard_incorrect_files.py

@@ -1,9 +1,7 @@
 import json
 import os
-from tqdm import tqdm
 
-import torch
-from torchvision.io import ImageReadMode, read_image
+from tqdm import tqdm
 
 JOINT_JSON_DIRECTORY = f"/home/{os.environ['USER']}/data/wit/all_jsons"
 SCALE_CONVERTED_DIRECTORY = f"/home/{os.environ['USER']}/data/wit_scale_converted"
@@ -16,13 +14,16 @@ for split in ["train", "valid", "test"]:
 
     supported_examples = []
     for example in tqdm(examples):
-        directory, filename = os.path.split(example['image_path'])
+        directory, filename = os.path.split(example["image_path"])
         if filename in valid_files:
             example["image_path"] = os.path.join(SCALE_CONVERTED_DIRECTORY, filename)
             supported_examples.append(json.dumps(example, ensure_ascii=False))
 
     print(f"Total {split} examples: {len(supported_examples)}")
-    with open(f"{SCALE_CONVERTED_DIRECTORY}/{split}_dataset_scale_converted_98_1_1_split.json", "w") as f:
+    with open(
+        f"{SCALE_CONVERTED_DIRECTORY}/{split}_dataset_scale_converted_98_1_1_split.json",
+        "w",
+    ) as f:
         f.write("\n".join(supported_examples))
 
 print("DONE!")

join_datasets_custom_split.py

@@ -1,10 +1,7 @@
-import os
 import json
+import os
 import random
 
-import pandas as pd
-
-
 DATA_DIR = f"/home/{os.environ['USER']}/data/wit/all_jsons"
 SEED = 0
 PROPORTION_TRAIN = 0.98
@@ -12,7 +9,9 @@ PROPORTION_VALID = 0.01
 
 random.seed(SEED)
 
-all_files = [f"{DATA_DIR}/{file_}" for file_ in os.listdir(DATA_DIR) if ("all" not in file_)]
+all_files = [
+    f"{DATA_DIR}/{file_}" for file_ in os.listdir(DATA_DIR) if ("all" not in file_)
+]
 
 print(all_files)
 
@@ -20,7 +19,9 @@ examples = []
 for file_ in all_files:
     print(file_)
     with open(file_) as f:
-        file_examples = [json.dumps(json.loads(line), ensure_ascii=False) for line in f.readlines()]
+        file_examples = [
+            json.dumps(json.loads(line), ensure_ascii=False) for line in f.readlines()
+        ]
     print(len(file_examples))
     examples.extend(file_examples)
 
@@ -34,15 +35,23 @@ random.shuffle(examples)
 print(examples[0])
 
 split_dataset = {}
-split_dataset["train"] = examples[:int(len(examples) * PROPORTION_TRAIN)]
-split_dataset["valid"] = examples[int(len(examples) * PROPORTION_TRAIN): int(len(examples) * (PROPORTION_TRAIN + PROPORTION_VALID))]
-split_dataset["test"] = examples[int(len(examples) * (PROPORTION_TRAIN + PROPORTION_VALID)):]
+split_dataset["train"] = examples[: int(len(examples) * PROPORTION_TRAIN)]
+split_dataset["valid"] = examples[
+    int(len(examples) * PROPORTION_TRAIN) : int(
+        len(examples) * (PROPORTION_TRAIN + PROPORTION_VALID)
+    )
+]
+split_dataset["test"] = examples[
+    int(len(examples) * (PROPORTION_TRAIN + PROPORTION_VALID)) :
+]
 
 
 for split in ["train", "valid", "test"]:
     print("-----")
     print(len(split_dataset[split]))
     print("-----")
-    with open(f"/home/{os.environ['USER']}/data/wit/all_jsons/{split}_dataset_all_98_1_1_split.json", "w") as f:
+    with open(
+        f"/home/{os.environ['USER']}/data/wit/all_jsons/{split}_dataset_all_98_1_1_split.json",
+        "w",
+    ) as f:
         f.write("\n".join(split_dataset[split]))
-

modeling_hybrid_clip.py

@@ -18,13 +18,13 @@ from typing import Optional, Tuple
 import flax.linen as nn
 import jax
 import jax.numpy as jnp
-from configuration_hybrid_clip import HybridCLIPConfig
 from flax.core.frozen_dict import FrozenDict
 from transformers import FLAX_MODEL_MAPPING, FlaxCLIPVisionModel
 from transformers.modeling_flax_utils import FlaxPreTrainedModel
 from transformers.models.clip.modeling_flax_clip import FlaxCLIPOutput
 from transformers.utils import logging
 
+from configuration_hybrid_clip import HybridCLIPConfig
 
 logger = logging.get_logger(__name__)
 
@@ -42,7 +42,9 @@ class FlaxHybridCLIPModule(nn.Module):
         self.vision_embed_dim = vision_config.hidden_size
 
         text_module = FLAX_MODEL_MAPPING[self.config.text_config.__class__].module_class
-        vision_module = FLAX_MODEL_MAPPING.get(self.config.vision_config.__class__, FlaxCLIPVisionModel).module_class
+        vision_module = FLAX_MODEL_MAPPING.get(
+            self.config.vision_config.__class__, FlaxCLIPVisionModel
+        ).module_class
 
         self.text_model = text_module(text_config, dtype=self.dtype)
         self.vision_model = vision_module(vision_config, dtype=self.dtype)
@@ -73,7 +75,9 @@ class FlaxHybridCLIPModule(nn.Module):
         output_hidden_states=None,
         return_dict=None,
     ):
-        return_dict = return_dict if return_dict is not None else self.config.return_dict
+        return_dict = (
+            return_dict if return_dict is not None else self.config.return_dict
+        )
 
         vision_outputs = self.vision_model(
             pixel_values=pixel_values,
@@ -101,7 +105,9 @@ class FlaxHybridCLIPModule(nn.Module):
         text_embeds = self.text_projection(text_embeds)
 
         # normalized features
-        image_embeds = image_embeds / jnp.linalg.norm(image_embeds, axis=-1, keepdims=True)
+        image_embeds = image_embeds / jnp.linalg.norm(
+            image_embeds, axis=-1, keepdims=True
+        )
         text_embeds = text_embeds / jnp.linalg.norm(text_embeds, axis=-1, keepdims=True)
 
         # cosine similarity as logits
@@ -110,7 +116,14 @@ class FlaxHybridCLIPModule(nn.Module):
         logits_per_image = logits_per_text.T
 
         if not return_dict:
-            return (logits_per_image, logits_per_text, text_embeds, image_embeds, text_outputs, vision_outputs)
+            return (
+                logits_per_image,
+                logits_per_text,
+                text_embeds,
+                image_embeds,
+                text_outputs,
+                vision_outputs,
+            )
 
         return FlaxCLIPOutput(
             logits_per_image=logits_per_image,
@@ -132,18 +145,30 @@ class FlaxHybridCLIP(FlaxPreTrainedModel):
         input_shape: Optional[Tuple] = None,
         seed: int = 0,
         dtype: jnp.dtype = jnp.float32,
-        **kwargs
+        **kwargs,
     ):
         if input_shape is None:
-            input_shape = ((1, 1), (1, config.vision_config.image_size, config.vision_config.image_size, 3))
+            input_shape = (
+                (1, 1),
+                (
+                    1,
+                    config.vision_config.image_size,
+                    config.vision_config.image_size,
+                    3,
+                ),
+            )
 
         module = self.module_class(config=config, dtype=dtype, **kwargs)
-        super().__init__(config, module, input_shape=input_shape, seed=seed, dtype=dtype)
+        super().__init__(
+            config, module, input_shape=input_shape, seed=seed, dtype=dtype
+        )
 
     def init_weights(self, rng: jax.random.PRNGKey, input_shape: Tuple) -> FrozenDict:
         # init input tensor
         input_ids = jnp.zeros(input_shape[0], dtype="i4")
-        position_ids = jnp.broadcast_to(jnp.arange(jnp.atleast_2d(input_ids).shape[-1]), input_shape[0])
+        position_ids = jnp.broadcast_to(
+            jnp.arange(jnp.atleast_2d(input_ids).shape[-1]), input_shape[0]
+        )
         token_type_ids = jnp.ones_like(input_ids)
         attention_mask = jnp.ones_like(input_ids)
 
@@ -152,7 +177,9 @@ class FlaxHybridCLIP(FlaxPreTrainedModel):
         params_rng, dropout_rng = jax.random.split(rng)
         rngs = {"params": params_rng, "dropout": dropout_rng}
 
-        return self.module.init(rngs, input_ids, pixel_values, attention_mask, position_ids, token_type_ids)["params"]
+        return self.module.init(
+            rngs, input_ids, pixel_values, attention_mask, position_ids, token_type_ids
+        )["params"]
 
     def __call__(
         self,
@@ -168,14 +195,24 @@ class FlaxHybridCLIP(FlaxPreTrainedModel):
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
     ):
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_attentions = (
+            output_attentions
+            if output_attentions is not None
+            else self.config.output_attentions
+        )
         output_hidden_states = (
-            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
+        return_dict = (
+            return_dict if return_dict is not None else self.config.return_dict
         )
-        return_dict = return_dict if return_dict is not None else self.config.return_dict
 
         if position_ids is None:
-            position_ids = jnp.broadcast_to(jnp.arange(jnp.atleast_2d(input_ids).shape[-1]), input_ids.shape)
+            position_ids = jnp.broadcast_to(
+                jnp.arange(jnp.atleast_2d(input_ids).shape[-1]), input_ids.shape
+            )
 
         if token_type_ids is None:
             token_type_ids = jnp.zeros_like(input_ids)
@@ -225,7 +262,9 @@ class FlaxHybridCLIP(FlaxPreTrainedModel):
             obtained by applying the projection layer to the pooled output of text model.
         """
         if position_ids is None:
-            position_ids = jnp.broadcast_to(jnp.arange(jnp.atleast_2d(input_ids).shape[-1]), input_ids.shape)
+            position_ids = jnp.broadcast_to(
+                jnp.arange(jnp.atleast_2d(input_ids).shape[-1]), input_ids.shape
+            )
 
         if token_type_ids is None:
             token_type_ids = jnp.zeros_like(input_ids)
@@ -238,7 +277,14 @@ class FlaxHybridCLIP(FlaxPreTrainedModel):
         if dropout_rng is not None:
             rngs["dropout"] = dropout_rng
 
-        def _get_features(module, input_ids, attention_mask, position_ids, token_type_ids, deterministic):
+        def _get_features(
+            module,
+            input_ids,
+            attention_mask,
+            position_ids,
+            token_type_ids,
+            deterministic,
+        ):
             text_outputs = module.text_model(
                 input_ids=input_ids,
                 attention_mask=attention_mask,
@@ -261,7 +307,9 @@ class FlaxHybridCLIP(FlaxPreTrainedModel):
             rngs=rngs,
         )
 
-    def get_image_features(self, pixel_values, dropout_rng: jax.random.PRNGKey = None, train=False):
+    def get_image_features(
+        self, pixel_values, dropout_rng: jax.random.PRNGKey = None, train=False
+    ):
         r"""
         Args:
             pixel_values (:obj:`numpy.ndarray` of shape :obj:`(batch_size, num_channels, height, width)`):
@@ -279,7 +327,9 @@ class FlaxHybridCLIP(FlaxPreTrainedModel):
             rngs["dropout"] = dropout_rng
 
         def _get_features(module, pixel_values, deterministic):
-            vision_outputs = module.vision_model(pixel_values=pixel_values, deterministic=deterministic)
+            vision_outputs = module.vision_model(
+                pixel_values=pixel_values, deterministic=deterministic
+            )
             pooled_output = vision_outputs[1]  # pooled_output
             image_features = module.visual_projection(pooled_output)
             return image_features
@@ -345,11 +395,15 @@ class FlaxHybridCLIP(FlaxPreTrainedModel):
         """
 
         kwargs_text = {
-            argument[len("text_") :]: value for argument, value in kwargs.items() if argument.startswith("text_")
+            argument[len("text_") :]: value
+            for argument, value in kwargs.items()
+            if argument.startswith("text_")
         }
 
         kwargs_vision = {
-            argument[len("vision_") :]: value for argument, value in kwargs.items() if argument.startswith("vision_")
+            argument[len("vision_") :]: value
+            for argument, value in kwargs.items()
+            if argument.startswith("vision_")
         }
 
         # remove text, vision kwargs from kwargs
@@ -372,7 +426,9 @@ class FlaxHybridCLIP(FlaxPreTrainedModel):
                 text_config = AutoConfig.from_pretrained(text_model_name_or_path)
                 kwargs_text["config"] = text_config
 
-            text_model = FlaxAutoModel.from_pretrained(text_model_name_or_path, *model_args, **kwargs_text)
+            text_model = FlaxAutoModel.from_pretrained(
+                text_model_name_or_path, *model_args, **kwargs_text
+            )
 
         vision_model = kwargs_vision.pop("model", None)
         if vision_model is None:
@@ -387,21 +443,29 @@ class FlaxHybridCLIP(FlaxPreTrainedModel):
                 vision_config = AutoConfig.from_pretrained(vision_model_name_or_path)
                 kwargs_vision["config"] = vision_config
 
-            vision_model = FlaxAutoModel.from_pretrained(vision_model_name_or_path, *model_args, **kwargs_vision)
+            vision_model = FlaxAutoModel.from_pretrained(
+                vision_model_name_or_path, *model_args, **kwargs_vision
+            )
 
         # instantiate config with corresponding kwargs
         dtype = kwargs.pop("dtype", jnp.float32)
-        config = HybridCLIPConfig.from_text_vision_configs(text_model.config, vision_model.config, **kwargs)
+        config = HybridCLIPConfig.from_text_vision_configs(
+            text_model.config, vision_model.config, **kwargs
+        )
 
         # init model
         model = cls(config, *model_args, dtype=dtype, **kwargs)
 
         if vision_config.model_type == "clip":
-            model.params["vision_model"]["vision_model"] = vision_model.params["vision_model"]
-            model.params["visual_projection"]["kernel"] = vision_model.params["visual_projection"]["kernel"]
+            model.params["vision_model"]["vision_model"] = vision_model.params[
+                "vision_model"
+            ]
+            model.params["visual_projection"]["kernel"] = vision_model.params[
+                "visual_projection"
+            ]["kernel"]
         else:
             model.params["vision_model"] = vision_model.params
 
         model.params["text_model"] = text_model.params
 
-        return model
+        return model

prepare_wit.py

@@ -3,14 +3,13 @@ import json
 import logging
 import os
 import time
-from typing import List
-import urllib.request
 import urllib.error
+import urllib.request
+from typing import List
 
 import pandas as pd
 from tqdm import tqdm
 
-
 logging.basicConfig(
     format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
     datefmt="%m/%d/%Y %H:%M:%S",
@@ -18,11 +17,18 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
-def split_and_save_datasets(lines: List[str], output_dir: str, train_proportion: float, valid_proportion: float):
+
+def split_and_save_datasets(
+    lines: List[str], output_dir: str, train_proportion: float, valid_proportion: float
+):
     total_lines = len(lines)
-    train_lines = lines[:int(total_lines * train_proportion)]
-    valid_lines = lines[int(total_lines * train_proportion):int(total_lines * (train_proportion + valid_proportion))]
-    test_lines = lines[int(total_lines * (train_proportion + valid_proportion)):]
+    train_lines = lines[: int(total_lines * train_proportion)]
+    valid_lines = lines[
+        int(total_lines * train_proportion) : int(
+            total_lines * (train_proportion + valid_proportion)
+        )
+    ]
+    test_lines = lines[int(total_lines * (train_proportion + valid_proportion)) :]
 
     with open(f"{output_dir}/train_dataset.json", "w") as f:
         f.write("\n".join(train_lines))
@@ -33,14 +39,33 @@ def split_and_save_datasets(lines: List[str], output_dir: str, train_proportion:
     with open(f"{output_dir}/test_dataset.json", "w") as f:
         f.write("\n".join(test_lines))
 
+
 def prepare_wit(
-    tsv: str, language: str, output_dir: str, seed: int, train_proportion: float, valid_proportion: float, backup_period: int, language_col: str="language", caption_col: str="caption_reference_description", url_col: str="image_url", pause=0.875, retries: int=10):
+    tsv: str,
+    language: str,
+    output_dir: str,
+    seed: int,
+    train_proportion: float,
+    valid_proportion: float,
+    backup_period: int,
+    language_col: str = "language",
+    caption_col: str = "caption_reference_description",
+    url_col: str = "image_url",
+    pause=0.875,
+    retries: int = 10,
+):
     os.makedirs(output_dir, exist_ok=True)
     logger.info("Loading dataset")
     df = pd.read_csv(tsv, sep="\t", engine="python")
     existing_files = set(os.listdir(output_dir))
-    not_exists_condition = (~(df[url_col].map(lambda x: x.split("/")[-1][-100:]).isin(existing_files)))
-    df = df[(df["language"] == language) & (~df["caption_reference_description"].isnull()) & not_exists_condition]
+    not_exists_condition = ~(
+        df[url_col].map(lambda x: x.split("/")[-1][-100:]).isin(existing_files)
+    )
+    df = df[
+        (df["language"] == language)
+        & (~df["caption_reference_description"].isnull())
+        & not_exists_condition
+    ]
     # Shuffle
     df = df.sample(frac=1.0, random_state=seed)
     logger.info(f"Trying to downloading {df.shape[0]} files")
@@ -58,14 +83,21 @@ def prepare_wit(
                     try:
                         # Download file
                         urllib.request.urlretrieve(url, image_path)
-                        lines.append(json.dumps({"image_path": image_path, "captions": [caption]}, ensure_ascii=False))
+                        lines.append(
+                            json.dumps(
+                                {"image_path": image_path, "captions": [caption]},
+                                ensure_ascii=False,
+                            )
+                        )
                         count += 1
                         break
-                    except urllib.error.HTTPError as e:
+                    except urllib.error.HTTPError:
                         time.sleep(pause * 10)
                 if count % backup_period == 0:
                     logger.info(f"Saving dataset backup: Number of lines {len(lines)}")
-                    split_and_save_datasets(lines, output_dir, train_proportion, valid_proportion)
+                    split_and_save_datasets(
+                        lines, output_dir, train_proportion, valid_proportion
+                    )
                 if retry == retries - 1:
                     logger.info(f"Skipping {image_filename}")
                 pbar.update(1)
@@ -73,16 +105,35 @@ def prepare_wit(
     finally:
         split_and_save_datasets(lines, output_dir, train_proportion, valid_proportion)
 
+
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description = "Download and prepare the WIT dataset")
-    parser.add_argument("--tsv", type=str, default=f"/home/{os.environ['USER']}/data/wit/wit_v1.train.all-1percent_sample.tsv")
+    parser = argparse.ArgumentParser(description="Download and prepare the WIT dataset")
+    parser.add_argument(
+        "--tsv",
+        type=str,
+        default=f"/home/{os.environ['USER']}/data/wit/wit_v1.train.all-1percent_sample.tsv",
+    )
     parser.add_argument("--language", type=str, default="es")
-    parser.add_argument("--output_dir", type=str, default=f"/home/{os.environ['USER']}/data/wit/prepared_dataset")
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default=f"/home/{os.environ['USER']}/data/wit/prepared_dataset",
+    )
     parser.add_argument("--random_seed", type=int, default=0)
     parser.add_argument("--train_proportion", type=float, default=0.8)
     parser.add_argument("--valid_proportion", type=float, default=0.1)
     parser.add_argument("--backup_period", type=int, default=1000)
 
     args = parser.parse_args()
-    assert args.train_proportion + args.valid_proportion < 1.0, "The sum of train_proportion and valid_proportion has to be < 1.0"
-    prepare_wit(args.tsv, args.language, args.output_dir, args.random_seed, args.train_proportion, args.valid_proportion, args.backup_period)
+    assert (
+        args.train_proportion + args.valid_proportion < 1.0
+    ), "The sum of train_proportion and valid_proportion has to be < 1.0"
+    prepare_wit(
+        args.tsv,
+        args.language,
+        args.output_dir,
+        args.random_seed,
+        args.train_proportion,
+        args.valid_proportion,
+        args.backup_period,
+    )

run_hybrid_clip.py

@@ -32,25 +32,26 @@ from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Callable, Optional
 
-import numpy as np
-import torch
-from torchvision.datasets import VisionDataset
-from torchvision.io import ImageReadMode, read_image
-from torchvision.transforms import CenterCrop, ConvertImageDtype, Normalize, Resize
-from torchvision.transforms.functional import InterpolationMode
-from tqdm import tqdm
-
 import jax
 import jax.numpy as jnp
+import numpy as np
 import optax
+import torch
 import transformers
 from flax import jax_utils
 from flax.jax_utils import unreplicate
 from flax.training import train_state
 from flax.training.common_utils import get_metrics, shard, shard_prng_key
-from modeling_hybrid_clip import FlaxHybridCLIP
-from transformers import AutoTokenizer, HfArgumentParser, TrainingArguments, is_tensorboard_available, set_seed
+from torchvision.datasets import VisionDataset
+from torchvision.io import ImageReadMode, read_image
+from torchvision.transforms import (CenterCrop, ConvertImageDtype, Normalize,
+                                    Resize)
+from torchvision.transforms.functional import InterpolationMode
+from tqdm import tqdm
+from transformers import (AutoTokenizer, HfArgumentParser, TrainingArguments,
+                          is_tensorboard_available, set_seed)
 
+from modeling_hybrid_clip import FlaxHybridCLIP
 
 logger = logging.getLogger(__name__)
 
@@ -61,7 +62,9 @@ if has_tensorboard:
         from flax.metrics.tensorboard import SummaryWriter
     except ImportError as ie:
         has_tensorboard = False
-        print(f"Unable to display metrics through TensorBoard because some package are not installed: {ie}")
+        print(
+            f"Unable to display metrics through TensorBoard because some package are not installed: {ie}"
+        )
 
 else:
     print(
@@ -90,20 +93,33 @@ class ModelArguments:
     )
     from_pt: bool = field(
         default=True,
-        metadata={"help": "whether to load the text and vision model using PyTorch checkpoints."},
+        metadata={
+            "help": "whether to load the text and vision model using PyTorch checkpoints."
+        },
     )
     config_name: Optional[str] = field(
-        default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
+        default=None,
+        metadata={
+            "help": "Pretrained config name or path if not the same as model_name"
+        },
     )
     tokenizer_name: Optional[str] = field(
-        default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}
+        default=None,
+        metadata={
+            "help": "Pretrained tokenizer name or path if not the same as model_name"
+        },
     )
     cache_dir: Optional[str] = field(
-        default=None, metadata={"help": "Where do you want to store the pretrained models downloaded from s3"}
+        default=None,
+        metadata={
+            "help": "Where do you want to store the pretrained models downloaded from s3"
+        },
     )
     use_fast_tokenizer: bool = field(
         default=True,
-        metadata={"help": "Whether to use one of the fast tokenizer (backed by the tokenizers library) or not."},
+        metadata={
+            "help": "Whether to use one of the fast tokenizer (backed by the tokenizers library) or not."
+        },
     )
     dtype: Optional[str] = field(
         default="float32",
@@ -119,9 +135,12 @@ class DataTrainingArguments:
     Arguments pertaining to what data we are going to input our model for training and eval.
     """
 
-    data_dir: Optional[str] = field(default=None, metadata={"help": "The data directory containing input files."})
+    data_dir: Optional[str] = field(
+        default=None, metadata={"help": "The data directory containing input files."}
+    )
     train_file: Optional[str] = field(
-        default=None, metadata={"help": "The input training data file (a jsonlines file)."}
+        default=None,
+        metadata={"help": "The input training data file (a jsonlines file)."},
     )
     validation_file: Optional[str] = field(
         default=None,
@@ -149,10 +168,12 @@ class DataTrainingArguments:
         },
     )
     overwrite_cache: bool = field(
-        default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
+        default=False,
+        metadata={"help": "Overwrite the cached training and evaluation sets"},
     )
     overwrite_cache: bool = field(
-        default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
+        default=False,
+        metadata={"help": "Overwrite the cached training and evaluation sets"},
     )
     preprocessing_num_workers: Optional[int] = field(
         default=None,
@@ -161,7 +182,9 @@ class DataTrainingArguments:
 
     def __post_init__(self):
         if self.train_file is None and self.validation_file is None:
-            raise ValueError("Need either a dataset name or a training/validation file.")
+            raise ValueError(
+                "Need either a dataset name or a training/validation file."
+            )
         else:
             if self.train_file is not None:
                 extension = self.train_file.split(".")[-1]
@@ -180,7 +203,10 @@ class Transform(torch.nn.Module):
             Resize([image_size], interpolation=InterpolationMode.BICUBIC),
             CenterCrop(image_size),
             ConvertImageDtype(torch.float),
-            Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
+            Normalize(
+                (0.48145466, 0.4578275, 0.40821073),
+                (0.26862954, 0.26130258, 0.27577711),
+            ),
         )
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -225,7 +251,7 @@ class ImageTextDataset(VisionDataset):
         self.image_paths = []
 
         for example in examples:
-            captions_subset =  example["captions"][:captions_per_image]
+            captions_subset = example["captions"][:captions_per_image]
             self.captions.extend(captions_subset)
             self.image_paths.extend([example["image_path"]] * len(captions_subset))
 
@@ -253,7 +279,9 @@ class TrainState(train_state.TrainState):
     dropout_rng: jnp.ndarray
 
     def replicate(self):
-        return jax_utils.replicate(self).replace(dropout_rng=shard_prng_key(self.dropout_rng))
+        return jax_utils.replicate(self).replace(
+            dropout_rng=shard_prng_key(self.dropout_rng)
+        )
 
 
 def write_metric(summary_writer, train_metrics, eval_metrics, train_time, step):
@@ -270,25 +298,39 @@ def write_metric(summary_writer, train_metrics, eval_metrics, train_time, step):
 
 
 def create_learning_rate_fn(
-    train_ds_size: int, train_batch_size: int, num_train_epochs: int, num_warmup_steps: int, learning_rate: float
+    train_ds_size: int,
+    train_batch_size: int,
+    num_train_epochs: int,
+    num_warmup_steps: int,
+    learning_rate: float,
 ) -> Callable[[int], jnp.array]:
     """Returns a linear warmup, linear_decay learning rate function."""
     steps_per_epoch = train_ds_size // train_batch_size
     num_train_steps = steps_per_epoch * num_train_epochs
-    warmup_fn = optax.linear_schedule(init_value=0.0, end_value=learning_rate, transition_steps=num_warmup_steps)
+    warmup_fn = optax.linear_schedule(
+        init_value=0.0, end_value=learning_rate, transition_steps=num_warmup_steps
+    )
     decay_fn = optax.linear_schedule(
-        init_value=learning_rate, end_value=0, transition_steps=num_train_steps - num_warmup_steps
+        init_value=learning_rate,
+        end_value=0,
+        transition_steps=num_train_steps - num_warmup_steps,
+    )
+    schedule_fn = optax.join_schedules(
+        schedules=[warmup_fn, decay_fn], boundaries=[num_warmup_steps]
     )
-    schedule_fn = optax.join_schedules(schedules=[warmup_fn, decay_fn], boundaries=[num_warmup_steps])
     return schedule_fn
 
 
 def main():
-    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))
+    parser = HfArgumentParser(
+        (ModelArguments, DataTrainingArguments, TrainingArguments)
+    )
     if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
         # If we pass only one argument to the script and it's the path to a json file,
         # let's parse it to get our arguments.
-        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
+        model_args, data_args, training_args = parser.parse_json_file(
+            json_file=os.path.abspath(sys.argv[1])
+        )
     else:
         model_args, data_args, training_args = parser.parse_args_into_dataclasses()
 
@@ -321,11 +363,15 @@ def main():
 
     if model_args.tokenizer_name:
         tokenizer = AutoTokenizer.from_pretrained(
-            model_args.tokenizer_name, cache_dir=model_args.cache_dir, use_fast=model_args.use_fast_tokenizer
+            model_args.tokenizer_name,
+            cache_dir=model_args.cache_dir,
+            use_fast=model_args.use_fast_tokenizer,
         )
     elif model_args.text_model_name_or_path:
         tokenizer = AutoTokenizer.from_pretrained(
-            model_args.text_model_name_or_path, cache_dir=model_args.cache_dir, use_fast=model_args.use_fast_tokenizer
+            model_args.text_model_name_or_path,
+            cache_dir=model_args.cache_dir,
+            use_fast=model_args.use_fast_tokenizer,
         )
     else:
         raise ValueError(
@@ -366,16 +412,28 @@ def main():
 
     # Store some constant
     num_epochs = int(training_args.num_train_epochs)
-    train_batch_size = int(training_args.per_device_train_batch_size) * jax.device_count()
+    train_batch_size = (
+        int(training_args.per_device_train_batch_size) * jax.device_count()
+    )
     eval_batch_size = int(training_args.per_device_eval_batch_size) * jax.device_count()
     steps_per_epoch = len(train_dataset) // train_batch_size
     total_train_steps = steps_per_epoch * num_epochs
 
     # Use collate function to tokenizer the text and convert the processed images to numpy
     def collate_fn(examples):
-        pixel_values = torch.stack([example[0] for example in examples]).permute(0, 2, 3, 1).numpy()
+        pixel_values = (
+            torch.stack([example[0] for example in examples])
+            .permute(0, 2, 3, 1)
+            .numpy()
+        )
         captions = [example[1] for example in examples]
-        inputs = tokenizer(captions, max_length=data_args.max_seq_length, padding="max_length", truncation=True, return_tensors="np")
+        inputs = tokenizer(
+            captions,
+            max_length=data_args.max_seq_length,
+            padding="max_length",
+            truncation=True,
+            return_tensors="np",
+        )
 
         batch = {
             "pixel_values": pixel_values,
@@ -408,7 +466,9 @@ def main():
 
     # Enable tensorboard only on the master node
     if has_tensorboard and jax.process_index() == 0:
-        summary_writer = SummaryWriter(log_dir=Path(training_args.output_dir).joinpath("logs").as_posix())
+        summary_writer = SummaryWriter(
+            log_dir=Path(training_args.output_dir).joinpath("logs").as_posix()
+        )
 
     # Initialize our training
     rng = jax.random.PRNGKey(training_args.seed)
@@ -433,7 +493,9 @@ def main():
     )
 
     # Setup train state
-    state = TrainState.create(apply_fn=model.__call__, params=model.params, tx=adamw, dropout_rng=dropout_rng)
+    state = TrainState.create(
+        apply_fn=model.__call__, params=model.params, tx=adamw, dropout_rng=dropout_rng
+    )
 
     def cross_entropy(logits, axis):
         logprobs = jax.nn.log_softmax(logits, axis=axis)
@@ -442,7 +504,9 @@ def main():
         return ce
 
     def clip_loss(similarity):
-        loss = (cross_entropy(similarity, axis=0) + cross_entropy(similarity, axis=1)) / 2
+        loss = (
+            cross_entropy(similarity, axis=0) + cross_entropy(similarity, axis=1)
+        ) / 2
         return loss
 
     # Define gradient update step fn
@@ -450,7 +514,9 @@ def main():
         dropout_rng, new_dropout_rng = jax.random.split(state.dropout_rng)
 
         def compute_loss(params):
-            logits = state.apply_fn(**batch, params=params, dropout_rng=dropout_rng, train=True)[0]
+            logits = state.apply_fn(
+                **batch, params=params, dropout_rng=dropout_rng, train=True
+            )[0]
             loss = clip_loss(logits)
             return loss
 
@@ -460,7 +526,10 @@ def main():
 
         new_state = state.apply_gradients(grads=grad, dropout_rng=new_dropout_rng)
 
-        metrics = {"loss": loss, "learning_rate": linear_decay_lr_schedule_fn(state.step)}
+        metrics = {
+            "loss": loss,
+            "learning_rate": linear_decay_lr_schedule_fn(state.step),
+        }
         metrics = jax.lax.pmean(metrics, axis_name="batch")
 
         return new_state, metrics
@@ -485,8 +554,12 @@ def main():
     logger.info("***** Running training *****")
     logger.info(f"  Num examples = {len(train_dataset)}")
     logger.info(f"  Num Epochs = {num_epochs}")
-    logger.info(f"  Instantaneous batch size per device = {training_args.per_device_train_batch_size}")
-    logger.info(f"  Total train batch size (w. parallel & distributed) = {train_batch_size}")
+    logger.info(
+        f"  Instantaneous batch size per device = {training_args.per_device_train_batch_size}"
+    )
+    logger.info(
+        f"  Total train batch size (w. parallel & distributed) = {train_batch_size}"
+    )
     logger.info(f"  Total optimization steps = {total_train_steps}")
 
     train_time = 0
@@ -504,7 +577,9 @@ def main():
         train_metrics = []
 
         steps_per_epoch = len(train_dataset) // train_batch_size
-        train_step_progress_bar = tqdm(total=steps_per_epoch, desc="Training...", position=1, leave=False)
+        train_step_progress_bar = tqdm(
+            total=steps_per_epoch, desc="Training...", position=1, leave=False
+        )
         # train
         for batch in train_loader:
             batch = shard(batch)
@@ -525,7 +600,9 @@ def main():
         # ======================== Evaluating ==============================
         eval_metrics = []
         eval_steps = len(eval_dataset) // eval_batch_size
-        eval_step_progress_bar = tqdm(total=eval_steps, desc="Evaluating...", position=2, leave=False)
+        eval_step_progress_bar = tqdm(
+            total=eval_steps, desc="Evaluating...", position=2, leave=False
+        )
         for batch in eval_loader:
             # Model forward
             batch = shard(batch)
@@ -541,14 +618,18 @@ def main():
 
         # Print metrics and update progress bar
         eval_step_progress_bar.close()
-        desc = f"Epoch... ({epoch + 1}/{num_epochs} | Eval Loss: {eval_metrics['loss']})"
+        desc = (
+            f"Epoch... ({epoch + 1}/{num_epochs} | Eval Loss: {eval_metrics['loss']})"
+        )
         epochs.write(desc)
         epochs.desc = desc
 
         # Save metrics
         if has_tensorboard and jax.process_index() == 0:
             cur_step = epoch * (len(train_dataset) // train_batch_size)
-            write_metric(summary_writer, train_metrics, eval_metrics, train_time, cur_step)
+            write_metric(
+                summary_writer, train_metrics, eval_metrics, train_time, cur_step
+            )
 
         # save checkpoint after each epoch and push checkpoint to the hub
         if jax.process_index() == 0:

scale_convert.py

@@ -1,12 +1,11 @@
 import glob
 import itertools
-from argparse import ArgumentParser
-from joblib import Parallel, delayed
 import os
 import subprocess
+from argparse import ArgumentParser
 from collections import Counter
-import shutil
 
+from joblib import Parallel, delayed
 
 parser = ArgumentParser()
 parser.add_argument("in_dir")
@@ -26,17 +25,16 @@ files = itertools.chain(
     glob.iglob(f"{args.in_dir}/*/*.SVG"),
 )
 
+
 def process_file(path):
     basename = os.path.basename(path)
-    ext = os.path.splitext(basename)[1]
     name = os.path.splitext(basename)[0]
 
-    dirname = os.path.dirname(path)
     try:
         r = subprocess.run(
             f'convert {path} -resize "224^>" -colorspace RGB -density 1200 {args.out_dir}/{name}.jpg',
             shell=True,
-            timeout=10
+            timeout=10,
         )
         rcode = r.returncode
     except subprocess.TimeoutExpired:
@@ -48,6 +46,8 @@ def process_file(path):
 
     return rcode
 
-codes = Parallel(n_jobs=32, prefer="threads", verbose=1)(delayed(process_file)(f) for f in files)
-print(Counter(codes))
 
+codes = Parallel(n_jobs=32, prefer="threads", verbose=1)(
+    delayed(process_file)(f) for f in files
+)
+print(Counter(codes))

test_on_image.py

@@ -17,13 +17,21 @@ def prepare_image(image_path, model):
     pixel_values = torch.stack([preprocessed_image]).permute(0, 2, 3, 1).numpy()
     return pixel_values
 
+
 def prepare_text(text, tokenizer):
     return tokenizer(text, return_tensors="np")
 
+
 def run_inference(image_path, text, model, tokenizer):
     pixel_values = prepare_image(image_path, model)
     input_text = prepare_text(text, tokenizer)
-    model_output = model(input_text["input_ids"], pixel_values, attention_mask=input_text["attention_mask"], train=False, return_dict=True)
+    model_output = model(
+        input_text["input_ids"],
+        pixel_values,
+        attention_mask=input_text["attention_mask"],
+        train=False,
+        return_dict=True,
+    )
     logits = model_output["logits_per_image"]
     score = jax.nn.sigmoid(logits)[0][0]
     return score
@@ -31,9 +39,11 @@ def run_inference(image_path, text, model, tokenizer):
 
 if __name__ == "__main__":
     model = FlaxHybridCLIP.from_pretrained("./")
-    tokenizer = AutoTokenizer.from_pretrained("bertin-project/bertin-roberta-base-spanish")
+    tokenizer = AutoTokenizer.from_pretrained(
+        "bertin-project/bertin-roberta-base-spanish"
+    )
 
     image_path = f"/home/{os.environ['USER']}/data/wit_scale_converted/Santuar.jpg"
     text = "Fachada del Santuario"
 
-    print(run_inference(image_path, text, model, tokenizer))
+    print(run_inference(image_path, text, model, tokenizer))