Add code to precompute embeddings.

app.py

@@ -47,6 +47,9 @@ def get_image_features():
 
 
 """
+
+# 👋 Ciao!
+
 # CLIP Italian Demo (Flax Community Week)
 """
 

precompute_features.py

@@ -0,0 +1,50 @@
+from tqdm import tqdm
+from argparse import ArgumentParser
+from jax import numpy as jnp
+from torchvision import datasets, transforms
+from torchvision.transforms import CenterCrop, Normalize, Resize, ToTensor
+from torchvision.transforms.functional import InterpolationMode
+from transformers import AutoTokenizer
+from modeling_hybrid_clip import FlaxHybridCLIP
+import utils
+import torch
+
+
+if __name__ == "__main__":
+    parser = ArgumentParser()
+    parser.add_argument("in_dir")
+    parser.add_argument("out_file")
+    args = parser.parse_args()
+
+    model = FlaxHybridCLIP.from_pretrained("clip-italian/clip-italian")
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        "dbmdz/bert-base-italian-xxl-uncased", cache_dir=None, use_fast=True
+    )
+
+    image_size = model.config.vision_config.image_size
+
+    val_preprocess = transforms.Compose(
+        [
+            Resize([image_size], interpolation=InterpolationMode.BICUBIC),
+            CenterCrop(image_size),
+            ToTensor(),
+            Normalize(
+                (0.48145466, 0.4578275, 0.40821073),
+                (0.26862954, 0.26130258, 0.27577711),
+            ),
+        ]
+    )
+
+    dataset = utils.CustomDataSet(args.in_dir, transform=val_preprocess)
+
+    loader = torch.utils.data.DataLoader(
+        dataset,
+        batch_size=256,
+        shuffle=False,
+        num_workers=16,
+        drop_last=False,
+    )
+
+    image_features = utils.precompute_image_features(model, loader)
+    jnp.save(f"static/features/{args.out_file}", image_features)