Merge pull request #112 from borisdayma/’cleanup’

.github/workflows/style.yml

@@ -0,0 +1,20 @@
+name: Lint
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+
+jobs:
+  lint:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v2
+      - uses: psf/black@stable
+      - uses: actions/setup-python@v2
+        with:
+          python-version: 3.9
+      - name: Install requirements
+        run: pip install ".[dev]"
+      - uses: jamescurtin/isort-action@master

Makefile

@@ -0,0 +1,5 @@
+.PHONY: style
+
+style:
+	black .
+	isort .

README.md

@@ -4,8 +4,8 @@ emoji: 🥑
 colorFrom: yellow
 colorTo: green
 sdk: streamlit
-app_file: app/app.py
-pinned: false
+app_file: app/streamlit/app.py
+pinned: True
 ---
 
 # DALL·E Mini
@@ -28,7 +28,9 @@ Refer to [our report](https://wandb.ai/dalle-mini/dalle-mini/reports/DALL-E-mini
 
 ### Dependencies Installation
 
-For development, use [`dev/requirements.txt`](dev/requirements.txt) or [`dev/environment.yaml`](dev/environment.yaml).
+For inference only, use `pip install git+https://github.com/borisdayma/dalle-mini.git`.
+
+For development, clone the repo and use `pip install -e ".[dev]"`.
 
 ### Training of VQGAN
 
@@ -42,15 +44,15 @@ Use [patil-suraj/vqgan-jax](https://github.com/patil-suraj/vqgan-jax).
 
 ### Training of Seq2Seq
 
-Refer to [`dev/seq2seq`](dev/seq2seq) folder.
+Use [`tools/train/train.py`](tools/train/train.py).
 
 You can also adjust the [sweep configuration file](https://docs.wandb.ai/guides/sweeps) if you need to perform a hyperparameter search.
 
 ### Inference Pipeline
 
-To generate sample predictions and understand the inference pipeline step by step, refer to [`dev/inference/inference_pipeline.ipynb`](dev/inference/inference_pipeline.ipynb).
+To generate sample predictions and understand the inference pipeline step by step, refer to [`tools/inference/inference_pipeline.ipynb`](tools/inference/inference_pipeline.ipynb).
 
-[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/borisdayma/dalle-mini/blob/main/dev/inference/inference_pipeline.ipynb)
+[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/borisdayma/dalle-mini/blob/main/tools/inference/inference_pipeline.ipynb)
 
 ## FAQ
 

app/gradio/app_gradio.py

@@ -2,51 +2,62 @@
 # coding: utf-8
 
 # Uncomment to run on cpu
-#import os
-#os.environ["JAX_PLATFORM_NAME"] = "cpu"
+# import os
+# os.environ["JAX_PLATFORM_NAME"] = "cpu"
 
 import random
 
+import gradio as gr
 import jax
-import flax.linen as nn
-from flax.training.common_utils import shard
-from flax.jax_utils import replicate, unreplicate
-
-from transformers import BartTokenizer, FlaxBartForConditionalGeneration
-
-from PIL import Image
 import numpy as np
-import matplotlib.pyplot as plt
+from flax.jax_utils import replicate
+from flax.training.common_utils import shard
+from PIL import Image, ImageDraw, ImageFont
 
+# ## CLIP Scoring
+from transformers import BartTokenizer, CLIPProcessor, FlaxCLIPModel
 from vqgan_jax.modeling_flax_vqgan import VQModel
+
 from dalle_mini.model import CustomFlaxBartForConditionalGeneration
 
-# ## CLIP Scoring
-from transformers import CLIPProcessor, FlaxCLIPModel
+DALLE_REPO = "flax-community/dalle-mini"
+DALLE_COMMIT_ID = "4d34126d0df8bc4a692ae933e3b902a1fa8b6114"
 
-import gradio as gr
+VQGAN_REPO = "flax-community/vqgan_f16_16384"
+VQGAN_COMMIT_ID = "90cc46addd2dd8f5be21586a9a23e1b95aa506a9"
 
-from dalle_mini.helpers import captioned_strip
+tokenizer = BartTokenizer.from_pretrained(DALLE_REPO, revision=DALLE_COMMIT_ID)
+model = CustomFlaxBartForConditionalGeneration.from_pretrained(
+    DALLE_REPO, revision=DALLE_COMMIT_ID
+)
+vqgan = VQModel.from_pretrained(VQGAN_REPO, revision=VQGAN_COMMIT_ID)
 
 
-DALLE_REPO = 'flax-community/dalle-mini'
-DALLE_COMMIT_ID = '4d34126d0df8bc4a692ae933e3b902a1fa8b6114'
+def captioned_strip(images, caption=None, rows=1):
+    increased_h = 0 if caption is None else 48
+    w, h = images[0].size[0], images[0].size[1]
+    img = Image.new("RGB", (len(images) * w // rows, h * rows + increased_h))
+    for i, img_ in enumerate(images):
+        img.paste(img_, (i // rows * w, increased_h + (i % rows) * h))
 
-VQGAN_REPO = 'flax-community/vqgan_f16_16384'
-VQGAN_COMMIT_ID = '90cc46addd2dd8f5be21586a9a23e1b95aa506a9'
+    if caption is not None:
+        draw = ImageDraw.Draw(img)
+        font = ImageFont.truetype(
+            "/usr/share/fonts/truetype/liberation2/LiberationMono-Bold.ttf", 40
+        )
+        draw.text((20, 3), caption, (255, 255, 255), font=font)
+    return img
 
-tokenizer = BartTokenizer.from_pretrained(DALLE_REPO, revision=DALLE_COMMIT_ID)
-model = CustomFlaxBartForConditionalGeneration.from_pretrained(DALLE_REPO, revision=DALLE_COMMIT_ID)
-vqgan = VQModel.from_pretrained(VQGAN_REPO, revision=VQGAN_COMMIT_ID)
 
 def custom_to_pil(x):
-    x = np.clip(x, 0., 1.)
-    x = (255*x).astype(np.uint8)
+    x = np.clip(x, 0.0, 1.0)
+    x = (255 * x).astype(np.uint8)
     x = Image.fromarray(x)
     if not x.mode == "RGB":
         x = x.convert("RGB")
     return x
 
+
 def generate(input, rng, params):
     return model.generate(
         **input,
@@ -59,9 +70,11 @@ def generate(input, rng, params):
         params=params,
     )
 
+
 def get_images(indices, params):
     return vqgan.decode_code(indices, params=params)
 
+
 p_generate = jax.pmap(generate, "batch")
 p_get_images = jax.pmap(get_images, "batch")
 
@@ -73,9 +86,16 @@ print("Initialize FlaxCLIPModel")
 processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
 print("Initialize CLIPProcessor")
 
+
 def hallucinate(prompt, num_images=64):
     prompt = [prompt] * jax.device_count()
-    inputs = tokenizer(prompt, return_tensors='jax', padding="max_length", truncation=True, max_length=128).data
+    inputs = tokenizer(
+        prompt,
+        return_tensors="jax",
+        padding="max_length",
+        truncation=True,
+        max_length=128,
+    ).data
     inputs = shard(inputs)
 
     all_images = []
@@ -92,6 +112,7 @@ def hallucinate(prompt, num_images=64):
             all_images.append(custom_to_pil(image))
     return all_images
 
+
 def clip_top_k(prompt, images, k=8):
     inputs = processor(text=prompt, images=images, return_tensors="np", padding=True)
     outputs = clip(**inputs)
@@ -99,24 +120,29 @@ def clip_top_k(prompt, images, k=8):
     scores = np.array(logits[0]).argsort()[-k:][::-1]
     return [images[score] for score in scores]
 
+
 def compose_predictions(images, caption=None):
     increased_h = 0 if caption is None else 48
     w, h = images[0].size[0], images[0].size[1]
-    img = Image.new("RGB", (len(images)*w, h + increased_h))
+    img = Image.new("RGB", (len(images) * w, h + increased_h))
     for i, img_ in enumerate(images):
-        img.paste(img_, (i*w, increased_h))
+        img.paste(img_, (i * w, increased_h))
 
     if caption is not None:
         draw = ImageDraw.Draw(img)
-        font = ImageFont.truetype("/usr/share/fonts/truetype/liberation2/LiberationMono-Bold.ttf", 40)
-        draw.text((20, 3), caption, (255,255,255), font=font)
+        font = ImageFont.truetype(
+            "/usr/share/fonts/truetype/liberation2/LiberationMono-Bold.ttf", 40
+        )
+        draw.text((20, 3), caption, (255, 255, 255), font=font)
     return img
 
+
 def top_k_predictions(prompt, num_candidates=32, k=8):
     images = hallucinate(prompt, num_images=num_candidates)
     images = clip_top_k(prompt, images, k=k)
     return images
 
+
 def run_inference(prompt, num_images=32, num_preds=8):
     images = top_k_predictions(prompt, num_candidates=num_images, k=num_preds)
     predictions = captioned_strip(images)
@@ -125,23 +151,28 @@ def run_inference(prompt, num_images=32, num_preds=8):
     """
     return (output_title, predictions)
 
+
 outputs = [
-    gr.outputs.HTML(label=""),      # To be used as title
-    gr.outputs.Image(label=''),
+    gr.outputs.HTML(label=""),  # To be used as title
+    gr.outputs.Image(label=""),
 ]
 
 description = """
 DALL·E-mini is an AI model that generates images from any prompt you give! Generate images from text:
 """
-gr.Interface(run_inference, 
-    inputs=[gr.inputs.Textbox(label='What do you want to see?')],
-    outputs=outputs, 
-    title='DALL·E mini',
+gr.Interface(
+    run_inference,
+    inputs=[gr.inputs.Textbox(label="What do you want to see?")],
+    outputs=outputs,
+    title="DALL·E mini",
     description=description,
     article="<p style='text-align: center'> Created by Boris Dayma et al. 2021 | <a href='https://github.com/borisdayma/dalle-mini'>GitHub</a> | <a href='https://wandb.ai/dalle-mini/dalle-mini/reports/DALL-E-mini--Vmlldzo4NjIxODA'>Report</a></p>",
-    layout='vertical',
-    theme='huggingface',
-    examples=[['an armchair in the shape of an avocado'], ['snowy mountains by the sea']],
+    layout="vertical",
+    theme="huggingface",
+    examples=[
+        ["an armchair in the shape of an avocado"],
+        ["snowy mountains by the sea"],
+    ],
     allow_flagging=False,
     live=False,
     # server_port=8999

app/gradio/app_gradio_ngrok.py

@@ -1,89 +0,0 @@
-#!/usr/bin/env python
-# coding: utf-8
-
-import requests
-from PIL import Image
-import numpy as np
-import matplotlib.pyplot as plt
-from io import BytesIO
-import base64
-import os
-
-import gradio as gr
-
-from dalle_mini.helpers import captioned_strip
-
-
-backend_url = os.environ["BACKEND_SERVER"]
-
-
-class ServiceError(Exception):
-    def __init__(self, status_code):
-        self.status_code = status_code
-
-def get_images_from_ngrok(prompt):
-    r = requests.post(
-        backend_url, 
-        json={"prompt": prompt}
-    )
-    if r.status_code == 200:
-        images = r.json()["images"]
-        images = [Image.open(BytesIO(base64.b64decode(img))) for img in images]
-        return images
-    else:
-        raise ServiceError(r.status_code)
-        
-def run_inference(prompt):
-    try:
-        images = get_images_from_ngrok(prompt)
-        predictions = captioned_strip(images)
-        output_title = f"""
-        <p style="font-size:22px; font-style:bold">Best predictions</p>
-        <p>We asked our model to generate 128 candidates for your prompt:</p>
-
-        <pre>
-
-        <b>{prompt}</b>
-        </pre>
-        <p>We then used a pre-trained <a href="https://huggingface.co/openai/clip-vit-base-patch32">CLIP model</a> to score them according to the
-        similarity of the text and the image representations.</p>
-
-        <p>This is the result:</p>
-        """
-        
-        output_description = """
-        <p>Read our <a style="color:blue;" href="https://wandb.ai/dalle-mini/dalle-mini/reports/DALL-E-mini--Vmlldzo4NjIxODA">full report</a> for more details on how this works.<p>
-        <p style='text-align: center'>Created with <a style="color:blue;" href="https://github.com/borisdayma/dalle-mini">DALL·E mini</a></p>
-        """
-
-    except ServiceError:
-        output_title = f"""
-        Sorry, there was an error retrieving the images. Please, try again later or <a href="mailto:pcuenca-dalle@guenever.net">contact us here</a>.
-        """
-        predictions = None
-        output_description = ""
-        
-    return (output_title, predictions, output_description)
-    
-outputs = [
-    gr.outputs.HTML(label=""),      # To be used as title
-    gr.outputs.Image(label=''),
-    gr.outputs.HTML(label=""),      # Additional text that appears in the screenshot
-]
-
-description = """
-Welcome to DALL·E-mini, a text-to-image generation model.
-"""
-gr.Interface(run_inference, 
-    inputs=[gr.inputs.Textbox(label='Prompt')],
-    outputs=outputs, 
-    title='DALL·E mini',
-    description=description,
-    article="<p style='text-align: center'> DALLE·mini by Boris Dayma et al. | <a href='https://github.com/borisdayma/dalle-mini'>GitHub</a></p>",
-    layout='vertical',
-    theme='huggingface',
-    examples=[['an armchair in the shape of an avocado'], ['snowy mountains by the sea']],
-    allow_flagging=False,
-    live=False,
-    # server_name="0.0.0.0",      # Bind to all interfaces
-).launch()

app/{app.py → streamlit/app.py}

@@ -1,9 +1,10 @@
 #!/usr/bin/env python
 # coding: utf-8
 
-from dalle_mini.backend import ServiceError, get_images_from_backend
 import streamlit as st
 
+from .backend import ServiceError, get_images_from_backend
+
 st.sidebar.markdown(
     """
 <style>
@@ -50,7 +51,7 @@ if prompt != "":
         <div class="st-b7">
         <div class="css-whx05o e13vu3m50">
         <div data-testid="stMarkdownContainer" class="css-1ekf893 e16nr0p30">
-                <img src="https://raw.githubusercontent.com/borisdayma/dalle-mini/main/app/img/loading.gif" width="30"/>
+                <img src="https://raw.githubusercontent.com/borisdayma/dalle-mini/main/app/streamlit/img/loading.gif" width="30"/>
                 Generating predictions for: <b>{prompt}</b>
         </div>
         </div>

{dalle_mini → app/streamlit}/backend.py

@@ -1,17 +1,17 @@
-import requests
-from io import BytesIO
 import base64
+from io import BytesIO
+
+import requests
 from PIL import Image
 
+
 class ServiceError(Exception):
     def __init__(self, status_code):
         self.status_code = status_code
 
+
 def get_images_from_backend(prompt, backend_url):
-    r = requests.post(
-        backend_url, 
-        json={"prompt": prompt}
-    )
+    r = requests.post(backend_url, json={"prompt": prompt})
     if r.status_code == 200:
         images = r.json()["images"]
         images = [Image.open(BytesIO(base64.b64decode(img))) for img in images]

dalle_mini/data.py

@@ -1,10 +1,12 @@
 from dataclasses import dataclass, field
-from datasets import load_dataset, Dataset
 from functools import partial
-import numpy as np
+
 import jax
 import jax.numpy as jnp
+import numpy as np
+from datasets import Dataset, load_dataset
 from flax.training.common_utils import shard
+
 from .text import TextNormalizer
 
 

dalle_mini/dataset.py

@@ -1,122 +0,0 @@
-"""
-An image-caption dataset dataloader.
-Luke Melas-Kyriazi, 2021
-"""
-import warnings
-from typing import Optional, Callable
-from pathlib import Path
-import numpy as np
-import torch
-import pandas as pd
-from torch.utils.data import Dataset
-from torchvision.datasets.folder import default_loader
-from PIL import ImageFile
-from PIL.Image import DecompressionBombWarning
-ImageFile.LOAD_TRUNCATED_IMAGES = True
-warnings.filterwarnings("ignore", category=UserWarning)
-warnings.filterwarnings("ignore", category=DecompressionBombWarning)
-
-
-class CaptionDataset(Dataset):
-    """
-    A PyTorch Dataset class for (image, texts) tasks. Note that this dataset 
-    returns the raw text rather than tokens. This is done on purpose, because
-    it's easy to tokenize a batch of text after loading it from this dataset.
-    """
-
-    def __init__(self, *, images_root: str, captions_path: str, text_transform: Optional[Callable] = None, 
-                 image_transform: Optional[Callable] = None, image_transform_type: str = 'torchvision',
-                 include_captions: bool = True):
-        """
-        :param images_root: folder where images are stored
-        :param captions_path: path to csv that maps image filenames to captions
-        :param image_transform: image transform pipeline
-        :param text_transform: image transform pipeline
-        :param image_transform_type: image transform type, either `torchvision` or `albumentations`
-        :param include_captions: Returns a dictionary with `image`, `text` if `true`; otherwise returns just the images.
-        """
-
-        # Base path for images
-        self.images_root = Path(images_root)
-
-        # Load captions as DataFrame
-        self.captions = pd.read_csv(captions_path, delimiter='\t', header=0)
-        self.captions['image_file'] = self.captions['image_file'].astype(str)
-
-        # PyTorch transformation pipeline for the image (normalizing, etc.)
-        self.text_transform = text_transform
-        self.image_transform = image_transform
-        self.image_transform_type = image_transform_type.lower()
-        assert self.image_transform_type in ['torchvision', 'albumentations']
-
-        # Total number of datapoints
-        self.size = len(self.captions)
-
-        # Return image+captions or just images
-        self.include_captions = include_captions
-
-    def verify_that_all_images_exist(self):
-        for image_file in self.captions['image_file']:
-            p = self.images_root / image_file
-            if not p.is_file():
-                print(f'file does not exist: {p}')
-
-    def _get_raw_image(self, i):
-        image_file = self.captions.iloc[i]['image_file']
-        image_path = self.images_root / image_file
-        image = default_loader(image_path)
-        return image
-
-    def _get_raw_text(self, i):
-        return self.captions.iloc[i]['caption']
-
-    def __getitem__(self, i):
-        image = self._get_raw_image(i)
-        caption = self._get_raw_text(i)
-        if self.image_transform is not None:
-            if self.image_transform_type == 'torchvision':
-                image = self.image_transform(image)
-            elif self.image_transform_type == 'albumentations':
-                image = self.image_transform(image=np.array(image))['image']
-            else:
-                raise NotImplementedError(f"{self.image_transform_type=}")
-        return {'image': image, 'text': caption} if self.include_captions else image
-
-    def __len__(self):
-        return self.size
-
-
-if __name__ == "__main__":
-    import albumentations as A
-    from albumentations.pytorch import ToTensorV2
-    from transformers import AutoTokenizer
-    
-    # Paths
-    images_root = './images'
-    captions_path = './images-list-clean.tsv'
-
-    # Create transforms
-    tokenizer = AutoTokenizer.from_pretrained('distilroberta-base')
-    def tokenize(text):
-        return tokenizer(text, max_length=32, truncation=True, return_tensors='pt', padding='max_length')
-    image_transform = A.Compose([
-        A.Resize(256, 256), A.CenterCrop(256, 256),
-        A.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]), ToTensorV2()])
-    
-    # Create dataset
-    dataset = CaptionDataset(
-        images_root=images_root,
-        captions_path=captions_path,
-        image_transform=image_transform,
-        text_transform=tokenize,
-        image_transform_type='albumentations')
-
-    # Create dataloader
-    dataloader = torch.utils.data.DataLoader(dataset, batch_size=2)
-    batch = next(iter(dataloader))
-    print({k: (v.shape if isinstance(v, torch.Tensor) else v) for k, v in batch.items()})
-
-    # # (Optional) Check that all the images exist
-    # dataset = CaptionDataset(images_root=images_root, captions_path=captions_path)
-    # dataset.verify_that_all_images_exist()
-    # print('Done')

dalle_mini/helpers.py

@@ -1,14 +0,0 @@
-from PIL import Image, ImageDraw, ImageFont
-
-def captioned_strip(images, caption=None, rows=1):
-    increased_h = 0 if caption is None else 48
-    w, h = images[0].size[0], images[0].size[1]
-    img = Image.new("RGB", (len(images)*w//rows, h*rows + increased_h))
-    for i, img_ in enumerate(images):
-        img.paste(img_, (i//rows*w, increased_h + (i % rows) * h))
-
-    if caption is not None:
-        draw = ImageDraw.Draw(img)
-        font = ImageFont.truetype("/usr/share/fonts/truetype/liberation2/LiberationMono-Bold.ttf", 40)
-        draw.text((20, 3), caption, (255,255,255), font=font)
-    return img

dalle_mini/model.py

@@ -1,16 +1,14 @@
-import jax
 import flax.linen as nn
-
+import jax
+from transformers import BartConfig
 from transformers.models.bart.modeling_flax_bart import (
-    FlaxBartModule,
-    FlaxBartForConditionalGenerationModule,
-    FlaxBartForConditionalGeneration,
-    FlaxBartEncoder,
     FlaxBartDecoder,
+    FlaxBartEncoder,
+    FlaxBartForConditionalGeneration,
+    FlaxBartForConditionalGenerationModule,
+    FlaxBartModule,
 )
 
-from transformers import BartConfig
-
 
 class CustomFlaxBartModule(FlaxBartModule):
     def setup(self):
@@ -46,6 +44,11 @@ class CustomFlaxBartForConditionalGenerationModule(
     FlaxBartForConditionalGenerationModule
 ):
     def setup(self):
+        # set default config
+        self.config.normalize_text = getattr(self.config, "normalize_text", False)
+        self.config.image_length = getattr(self.config, "image_length", 256)
+        self.config.image_vocab_size = getattr(self.config, "image_vocab_size", 16384)
+
         self.model = CustomFlaxBartModule(config=self.config, dtype=self.dtype)
         self.lm_head = nn.Dense(
             self.config.image_vocab_size + 1,  # encoded image token space + 1 for bos

dalle_mini/text.py

@@ -2,13 +2,15 @@
 Utilities for processing text.
 """
 
+import html
+import math
+import random
+import re
 from pathlib import Path
-from unidecode import unidecode
 
-import re, math, random, html
 import ftfy
-
 from huggingface_hub import hf_hub_download
+from unidecode import unidecode
 
 # based on wiki word occurence
 person_token = [("a person", 282265), ("someone", 121194), ("somebody", 12219)]

dev/README.md

@@ -1,122 +0,0 @@
-# Development Instructions for TPU
-
-## Setup
-
-- Apply to the [TRC program](https://sites.research.google/trc/) for free TPU credits if you're elligible.
-- Follow the [Cloud TPU VM User's Guide](https://cloud.google.com/tpu/docs/users-guide-tpu-vm) to set up gcloud.
-- Verify `gcloud config list`, in particular account, project & zone.
-- Create a TPU VM per the guide and connect to it.
-
-When needing a larger disk:
-
-- Create a balanced persistent disk (SSD, so pricier than default HDD but much faster): `gcloud compute disks create DISK_NAME --size SIZE_IN_GB --type pd-balanced`
-- Attach the disk to your instance by adding `--data-disk source=REF` per ["Adding a persistent disk to a TPU VM" guide](https://cloud.google.com/tpu/docs/setup-persistent-disk), eg `gcloud alpha compute tpus tpu-vm create INSTANCE_NAME --accelerator-type=v3-8 --version=v2-alpha --data-disk source=projects/tpu-toys/zones/europe-west4-a/disks/DISK_NAME`
-- Format the partition as described in the guide.
-- Make sure to set up automatic remount of disk at restart.
-
-## Connect VS Code
-
-- Find external IP in the UI or with `gcloud alpha compute tpus tpu-vm describe INSTANCE_NAME`
-- Verify you can connect in terminal with `ssh EXTERNAL_IP -i ~/.ssh/google_compute_engine`
-- Add the same command as ssh host in VS Code.
-- Check config file
-
-  ```
-  Host INSTANCE_NAME
-    HostName EXTERNAL_IP
-    IdentityFile ~/.ssh/google_compute_engine
-  ```
-
-## Environment configuration
-
-### Use virtual environments (optional)
-
-We recommend using virtual environments (such as conda, venv or pyenv-virtualenv).
-
-If you want to use `pyenv` and `pyenv-virtualenv`:
-
-- Installation
-
-  - [Set up build environment](https://github.com/pyenv/pyenv/wiki#suggested-build-environment)
-  - Use [pyenv-installer](https://github.com/pyenv/pyenv-installer): `curl https://pyenv.run | bash`
-  - bash set-up:
-
-    ```bash
-    echo '\n'\
-        '# pyenv setup \n'\
-        'export PYENV_ROOT="$HOME/.pyenv" \n'\
-        'export PATH="$PYENV_ROOT/bin:$PATH" \n'\
-        'eval "$(pyenv init --path)" \n'\
-        'eval "$(pyenv init -)" \n'\
-        'eval "$(pyenv virtualenv-init -)"' >> ~/.bashrc
-    ```
-
-- Usage
-
-  - Install a python version: `pyenv install X.X.X`
-  - Create a virtual environment: `pyenv virtualenv 3.9.6 dalle_env`
-  - Activate: `pyenv activate dalle_env`
-
-    Note: you can auto-activate your environment at a location with `echo dalle_env >> .python-version`
-
-### Tools
-
-- Git
-
-  - `git config --global user.email "name@domain.com"
-  - `git config --global user.name "First Last"
-
-- Github CLI
-
-  - See [installation instructions](https://github.com/cli/cli/blob/trunk/docs/install_linux.md)
-  - `gh auth login`
-
-- Direnv
-
-  - Install direnv: `sudo apt-get update && sudo apt-get install direnv`
-  - bash set-up:
-
-    ```bash
-    echo -e '\n'\
-        '# direnv setup \n'\
-        'eval "$(direnv hook bash)" \n' >> ~/.bashrc
-    ```
-
-### Set up repo
-
-- Clone repo: `gh repo clone borisdayma/dalle-mini`
-- If using `pyenv-virtualenv`, auto-activate env: `echo dalle_env >> .python-version`
-
-## Environment
-
-- Install the following (use it later to update our dev requirements.txt)
-
-```
-requests
-pillow
-jupyterlab
-ipywidgets
-
--e ../datasets[streaming]
--e ../transformers
--e ../webdataset
-
-# JAX
---find-links https://storage.googleapis.com/jax-releases/libtpu_releases.html
-jax[tpu]>=0.2.16
-flax
-```
-
-- `transformers-cli login`
-
----
-
-- set `HF_HOME="/mnt/disks/persist/cache/huggingface"` in `/etc/environment` and ensure you have required permissions, then restart.
-
-## Working with datasets or models
-
-- Install [Git LFS](https://github.com/git-lfs/git-lfs/wiki/Installation)
-- Clone a dataset without large files: `GIT_LFS_SKIP_SMUDGE=1 git clone https://huggingface.co/datasets/.../...`
-- Use a local [credential store](https://git-scm.com/book/en/v2/Git-Tools-Credential-Storage) for caching credentials
-- Track specific extentions: `git lfs track "*.ext"`
-- See files tracked with LFS with `git lfs ls-files`

dev/data/CC12M_downloader.py

@@ -1,91 +0,0 @@
-# Luke Melas-Kyriazi's code. (https://twitter.com/lukemelas)
-
-#%% 
-import sys
-import os
-from datetime import datetime
-import pandas as pd
-import contexttimer
-from urllib.request import urlopen
-import requests
-from PIL import Image
-import torch
-from torchvision.transforms import functional as TF
-from multiprocessing import Pool
-from tqdm import tqdm
-import logging
-
-# Setup
-logging.basicConfig(filename='download.log', filemode='w', level=logging.INFO)
-requests.packages.urllib3.disable_warnings(requests.packages.urllib3.exceptions.InsecureRequestWarning)
-
-
-# # For downloading SVG images (I can't get this to work)
-# from io import BytesIO
-# import cairosvg
-
-#%% 
-# Load data
-print(f'Starting to load at {datetime.now().isoformat(timespec="minutes")}')
-with contexttimer.Timer(prefix="Loading from tsv"):
-    df = pd.read_csv('./cc12m.tsv', delimiter='\t', header=None)
-
-url_to_idx_map = {url: index for index, url, caption in df.itertuples()}
-print(f'Loaded {len(url_to_idx_map)} urls')
-
-#%% 
-df.head()
-
-#%% 
-
-# Note: it seems that there are no SVG images
-df.sample(10000)[1].str.contains('.svg').sum()
-
-#%% 
-# Resize function
-def resize(img):
-    max_size_of_short_side = 512
-    if min(img.size) > max_size_of_short_side:
-        img = TF.resize(img, size=max_size_of_short_side, interpolation=Image.LANCZOS) 
-    return img
-
-base_dir = os.path.join(os.getcwd(), 'images')
-
-def process(item):
-    url, image_id = item
-    try:
-        base_url = os.path.basename(url)  # extract base url
-        stem, ext = os.path.splitext(base_url)  # split into stem and extension
-        filename = f'{image_id:08d}---{stem}.jpg'  # create filename
-        filepath = os.path.join(base_dir, filename)  # concat to get filepath
-        if not os.path.isfile(filepath):
-            # if filepath.endswith('.svg'):
-            #     raise NotImplementedError()
-            #     image_bytes = BytesIO()  # create a bytestream
-            #     cairosvg.svg2png(url=url, write_to=image_bytes)  # convert svg into image
-            # else:
-            req = requests.get(url, stream=True, timeout=1, verify=False).raw
-            image = Image.open(req).convert('RGB')
-            if min(image.size) > 512:
-                image = TF.resize(image, size=512, interpolation=Image.LANCZOS)
-            # image = resize(image)  # resize PIL image
-            image.save(filepath)  # save PIL image
-    except Exception as e:
-        logging.info(" ".join(repr(e).splitlines()))
-        logging.error(url)
-
-#%% 
-#for i, item in enumerate(tqdm(url_to_idx_map.items(), total=len(url_to_idx_map))):
-#    process(item)
-#    if i > 100:
-#        break
-
-# Use multiprocessing for speed
-list_of_items = list(url_to_idx_map.items())
-print(len(list_of_items))
-list_of_items = list_of_items[10_000_000:]
-print(len(list_of_items))
-with Pool(128) as p:
-    r = list(tqdm(p.imap(process, list_of_items), total=len(list_of_items)))
-    print('DONE')
-

dev/data/CC3M_downloader.py

@@ -1,62 +0,0 @@
-'''
-This script was adapted from Luke Melas-Kyriazi's code. (https://twitter.com/lukemelas)
-Few changes were made for the particular dataset. You're required to have the `.tsv` file downloaded in your directory.
-Find them here- [https://github.com/google-research-datasets/conceptual-captions]
-'''
-
-import sys
-import os
-from datetime import datetime
-import pandas as pd
-import contexttimer
-from urllib.request import urlopen
-import requests
-from PIL import Image
-import torch
-from torchvision.transforms import functional as TF
-from multiprocessing import Pool
-from tqdm import tqdm
-import logging
-import sys
-
-# Setup
-logging.basicConfig(filename='download.log', filemode='w', level=logging.INFO)
-requests.packages.urllib3.disable_warnings(requests.packages.urllib3.exceptions.InsecureRequestWarning)
-
-if len(sys.argv) != 3:
-    print("Provide .tsv file name & output directory. e.g. python downloader.py Train-GCC-training.tsv training")
-    exit(1)
-
-# Load data
-print(f'Starting to load at {datetime.now().isoformat(timespec="minutes")}')
-with contexttimer.Timer(prefix="Loading from tsv"):
-    df = pd.read_csv(sys.argv[1], delimiter='\t', header=None)
-
-url_to_idx_map = {url: index for index, caption, url in df.itertuples()}
-print(f'Loaded {len(url_to_idx_map)} urls')
-
-base_dir = os.path.join(os.getcwd(), sys.argv[2])
-
-def process(item):
-    url, image_id = item
-    try:
-        base_url = os.path.basename(url)  # extract base url
-        stem, ext = os.path.splitext(base_url)  # split into stem and extension
-        filename = f'{image_id:08d}---{stem}.jpg'  # create filename
-        filepath = os.path.join(base_dir, filename)  # concat to get filepath
-        if not os.path.isfile(filepath):
-            req = requests.get(url, stream=True, timeout=1, verify=False).raw
-            image = Image.open(req).convert('RGB')
-            if min(image.size) > 512:
-                image = TF.resize(image, size=512, interpolation=Image.LANCZOS)
-            image.save(filepath)  # save PIL image
-    except Exception as e:
-        logging.info(" ".join(repr(e).splitlines()))
-        logging.error(url)
-
-list_of_items = list(url_to_idx_map.items())
-print(len(list_of_items))
-
-with Pool(128) as p:
-    r = list(tqdm(p.imap(process, list_of_items), total=len(list_of_items)))
-    print('DONE')

dev/data/README.md

@@ -1,3 +0,0 @@
-# Data
-
-Utility scripts for downloading CC3M and CC12M.

dev/encoding/vqgan-jax-encoding-streaming.ipynb

@@ -1,562 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "markdown",
-   "id": "d0b72877",
-   "metadata": {},
-   "source": [
-    "# VQGAN JAX Encoding for 🤗 Datasets in streaming mode"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "ba7b31e6",
-   "metadata": {},
-   "source": [
-    "This notebook shows how to pre-encode images to token sequences using JAX, VQGAN and 🤗 Datasets in streaming mode.\n",
-    "\n",
-    "This example uses our YFCC100M dataset, but it should be easy to adapt to any other image/caption dataset in the huggingface hub."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "3b59489e",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import io\n",
-    "\n",
-    "import requests\n",
-    "from PIL import Image\n",
-    "import numpy as np\n",
-    "from tqdm import tqdm\n",
-    "\n",
-    "import torch\n",
-    "import torchvision.transforms as T\n",
-    "import torchvision.transforms.functional as TF\n",
-    "from torchvision.transforms import InterpolationMode\n",
-    "import os\n",
-    "\n",
-    "import jax\n",
-    "from jax import pmap"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "c7c4c1e6",
-   "metadata": {},
-   "source": [
-    "## Dataset and Parameters"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "d45a289e",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import datasets\n",
-    "from datasets import Dataset, load_dataset"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "f26e4f18",
-   "metadata": {},
-   "source": [
-    "We'll use the `validation` set for testing. Adjust accordingly."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "28893c3e",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "dataset = load_dataset('dalle-mini/YFCC100M_OpenAI_subset', use_auth_token=True, streaming=True, split='validation')"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "33861477",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from pathlib import Path\n",
-    "\n",
-    "yfcc100m = Path.home()/'data'/'YFCC100M_OpenAI_subset'\n",
-    "yfcc100m_output = yfcc100m/'encoded'      # Output directory for encoded files"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "6e7b71c4",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "batch_size = 128     # Per device\n",
-    "num_workers = 16     # Unused in streaming mode"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "0793c26a",
-   "metadata": {},
-   "source": [
-    "### Data preparation"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "86415769",
-   "metadata": {},
-   "source": [
-    "* Images: we transform them so they are center-cropped and square, all of the same size so we can build batches for TPU/GPU processing.\n",
-    "* Captions: we extract a single `caption` column from the source data, by concatenating the cleaned title and description.\n",
-    "\n",
-    "These transformations are done using the Datasets `map` function. In the case of streaming datasets, transformations will run as needed instead of pre-processing the dataset at once."
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "0fdf1851",
-   "metadata": {},
-   "source": [
-    "This helper function is used to decode images from the bytes retrieved in `streaming` mode."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "5bbca804",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from PIL import Image\n",
-    "import io\n",
-    "\n",
-    "def get_image(byte_stream):\n",
-    "    image = Image.open(io.BytesIO(byte_stream))\n",
-    "    return image.convert('RGB')"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "b435290b",
-   "metadata": {},
-   "source": [
-    "Image processing"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "7e73dfa3",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def center_crop(image, max_size=256):\n",
-    "    # Note: we allow upscaling too. We should exclude small images.    \n",
-    "    image = TF.resize(image, max_size, interpolation=InterpolationMode.LANCZOS)\n",
-    "    image = TF.center_crop(image, output_size=2 * [max_size])\n",
-    "    return image\n",
-    "\n",
-    "preprocess_image = T.Compose([\n",
-    "    get_image,\n",
-    "    center_crop,\n",
-    "    T.ToTensor(),\n",
-    "    lambda t: t.permute(1, 2, 0)   # Reorder, we need dimensions last\n",
-    "])"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "1e3ac8de",
-   "metadata": {},
-   "source": [
-    "Caption preparation"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "aadb4d23",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import string\n",
-    "\n",
-    "def create_caption(title, description):\n",
-    "    title = title.strip()\n",
-    "    description = description.strip()\n",
-    "    if len(title) > 0 and title[-1] not in '.!?': title += '.'\n",
-    "    return f'{title} {description}'"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "3c4522b9",
-   "metadata": {},
-   "source": [
-    "And this is the basic transformation function to use in `map`. We don't really need the `key`, but we'll keep it for reference. Since we are returning a new dictionary (as opposed to adding entries to the input), this also removes any metadata columns we don't need."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "2566ff68",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def prepare_item(item):\n",
-    "    return {\n",
-    "        'key': item['key'],\n",
-    "        'caption': create_caption(item['title_clean'], item['description_clean']),\n",
-    "        'image': preprocess_image(item['img'])\n",
-    "    }"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "e519e475",
-   "metadata": {},
-   "source": [
-    "Unlike when using non-streaming datasets, the following operation completes immediately in streaming mode. In streaming mode, `num_proc` is not supported."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "10d7750e",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "prepared_dataset = dataset.map(prepare_item, batched=False)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "a8595539",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "%%time\n",
-    "item = next(iter(prepared_dataset))"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "04a6eeb4",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "assert(list(item.keys()) == ['key', 'caption', 'image'])"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "40d3115f",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "item['image'].shape"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "dd844e1c",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "T.ToPILImage()(item['image'].permute(2, 0, 1))"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "44d50a51",
-   "metadata": {},
-   "source": [
-    "### Torch DataLoader"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "17a4bbc6",
-   "metadata": {},
-   "source": [
-    "We'll create a PyTorch DataLoader for convenience. This allows us to easily take batches of our desired size.\n",
-    "\n",
-    "We won't be using parallel processing of the DataLoader for now, as the items will be retrieved on the fly. We could attempt to do it using these recommendations: https://pytorch.org/docs/stable/data.html#multi-process-data-loading. For performance considerations, please refer to this thread: https://discuss.huggingface.co/t/allow-streaming-of-large-datasets-with-image-audio/8062/13"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "e1c08b7e",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import torch\n",
-    "from torch.utils.data import DataLoader"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "6a296677",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "torch_dataset = prepared_dataset.with_format(\"torch\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "29ab13bc",
-   "metadata": {},
-   "source": [
-    "**Note**: according to my tests, `num_workers` is not compatible with Datasets in streaming mode. Processes deadlock and there's no progress."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "e2df5e13",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "dataloader = DataLoader(torch_dataset, batch_size=batch_size * jax.device_count())"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "c15e3783",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "batch = next(iter(dataloader))"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "71d027fe",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "batch['image'].shape"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "a354472b",
-   "metadata": {},
-   "source": [
-    "## VQGAN-JAX model"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "2fcf01d7",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from vqgan_jax.modeling_flax_vqgan import VQModel"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "9daa636d",
-   "metadata": {},
-   "source": [
-    "We'll use a VQGAN trained with Taming Transformers and converted to a JAX model."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "47a8b818",
-   "metadata": {
-    "scrolled": true
-   },
-   "outputs": [],
-   "source": [
-    "model = VQModel.from_pretrained(\"flax-community/vqgan_f16_16384\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "62ad01c3",
-   "metadata": {},
-   "source": [
-    "## Encoding"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "20357f74",
-   "metadata": {},
-   "source": [
-    "Encoding is really simple using `shard` to automatically distribute \"superbatches\" across devices, and `pmap`. This is all it takes to create our encoding function, that will be jitted on first use."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "6686b004",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from flax.training.common_utils import shard\n",
-    "from functools import partial"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "322a4619",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "@partial(jax.pmap, axis_name=\"batch\")\n",
-    "def encode(batch):\n",
-    "    # Not sure if we should `replicate` params, does not seem to have any effect\n",
-    "    _, indices = model.encode(batch)\n",
-    "    return indices"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "14375a41",
-   "metadata": {},
-   "source": [
-    "### Encoding loop"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "ff6c10d4",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import os\n",
-    "import pandas as pd\n",
-    "\n",
-    "def encode_captioned_dataset(dataloader, output_dir, save_every=14):\n",
-    "    output_dir.mkdir(parents=True, exist_ok=True)\n",
-    "        \n",
-    "    # Saving strategy:\n",
-    "    # - Create a new file every so often to prevent excessive file seeking.\n",
-    "    # - Save each batch after processing.\n",
-    "    # - Keep the file open until we are done with it.\n",
-    "    file = None        \n",
-    "    for n, batch in enumerate(tqdm(iter(dataloader))):\n",
-    "        if (n % save_every == 0):\n",
-    "            if file is not None:\n",
-    "                file.close()\n",
-    "            split_num = n // save_every\n",
-    "            file = open(output_dir/f'split_{split_num:05x}.jsonl', 'w')\n",
-    "\n",
-    "        images = batch[\"image\"].numpy()\n",
-    "        images = shard(images.squeeze())\n",
-    "        encoded = encode(images)\n",
-    "        encoded = encoded.reshape(-1, encoded.shape[-1])\n",
-    "\n",
-    "        keys = batch[\"key\"]\n",
-    "        captions = batch[\"caption\"]\n",
-    "\n",
-    "        encoded_as_string = list(map(lambda item: np.array2string(item, separator=',', max_line_width=50000, formatter={'int':lambda x: str(x)}), encoded))\n",
-    "        batch_df = pd.DataFrame.from_dict({\"key\": keys, \"caption\": captions, \"encoding\": encoded_as_string})\n",
-    "        batch_df.to_json(file, orient='records', lines=True)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "09ff75a3",
-   "metadata": {},
-   "source": [
-    "Create a new file every 318 iterations. This should produce splits of ~500 MB each, when using a total batch size of 1024."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "96222bb4",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "save_every = 318"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "7704863d",
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "28it [01:17,  1.60s/it]"
-     ]
-    }
-   ],
-   "source": [
-    "encode_captioned_dataset(dataloader, yfcc100m_output, save_every=save_every)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "e266a70a",
-   "metadata": {},
-   "source": [
-    "This is ~10-15 slower than local encoding from an SSD. For performance considerations, see the discussion at https://discuss.huggingface.co/t/allow-streaming-of-large-datasets-with-image-audio/8062/13."
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "8953dd84",
-   "metadata": {},
-   "source": [
-    "----"
-   ]
-  }
- ],
- "metadata": {
-  "interpreter": {
-   "hash": "db471c52d602b4f5f40ecaf278e88ccfef85c29d0a1a07185b0d51fc7acf4e26"
-  },
-  "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.8.10"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 5
-}

dev/encoding/vqgan-jax-encoding-with-captions.ipynb

@@ -1,355 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "markdown",
-   "id": "d0b72877",
-   "metadata": {},
-   "source": [
-    "# vqgan-jax-encoding-with-captions"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "875c82b3",
-   "metadata": {},
-   "source": [
-    "Notebook based on [vqgan-jax-reconstruction](https://colab.research.google.com/drive/1mdXXsMbV6K_LTvCh3IImRsFIWcKU5m1w?usp=sharing) by @surajpatil.\n",
-    "\n",
-    "We process a `tsv` file with `image_file` and `caption` fields, and add a `vqgan_indices` column with indices extracted from a VQGAN-JAX model."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 1,
-   "id": "3b59489e",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import io\n",
-    "\n",
-    "import requests\n",
-    "from PIL import Image\n",
-    "import numpy as np\n",
-    "from tqdm import tqdm\n",
-    "\n",
-    "import torch\n",
-    "import torchvision.transforms as T\n",
-    "import torchvision.transforms.functional as TF\n",
-    "from torchvision.transforms import InterpolationMode\n",
-    "from torch.utils.data import Dataset, DataLoader\n",
-    "\n",
-    "import jax\n",
-    "from jax import pmap"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "511c3b9e",
-   "metadata": {},
-   "source": [
-    "## VQGAN-JAX model"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 2,
-   "id": "2ca50dc7",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from vqgan_jax.modeling_flax_vqgan import VQModel"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "7b60da9a",
-   "metadata": {},
-   "source": [
-    "We'll use a VQGAN trained by using Taming Transformers and converted to a JAX model."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 3,
-   "id": "29ce8b15",
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "db406bdfc5d5428eaeae1631a04989dd",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Downloading:   0%|          | 0.00/433 [00:00<?, ?B/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "3e37f07fba6d48fca70313ae1fa8cc32",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Downloading:   0%|          | 0.00/304M [00:00<?, ?B/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "INFO:absl:Starting the local TPU driver.\n",
-      "INFO:absl:Unable to initialize backend 'tpu_driver': Not found: Unable to find driver in registry given worker: local://\n",
-      "INFO:absl:Unable to initialize backend 'gpu': Not found: Could not find registered platform with name: \"cuda\". Available platform names are: Interpreter Host TPU\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Working with z of shape (1, 256, 16, 16) = 65536 dimensions.\n"
-     ]
-    }
-   ],
-   "source": [
-    "model = VQModel.from_pretrained(\"flax-community/vqgan_f16_16384\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "c7c4c1e6",
-   "metadata": {},
-   "source": [
-    "## Dataset"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "7014a7ce",
-   "metadata": {},
-   "source": [
-    "We use Luke Melas-Kyriazi's `dataset.py` which reads image paths and captions from a tsv file that contains both. We only need the images for encoding."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 4,
-   "id": "85832702",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from dalle_mini.dataset import *"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 5,
-   "id": "81b19eca",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "cc12m_images = '/data/CC12M/images'\n",
-    "cc12m_list = '/data/CC12M/images-list-clean.tsv'\n",
-    "# cc12m_list = '/data/CC12M/images-10000.tsv'\n",
-    "cc12m_output = '/data/CC12M/images-encoded.tsv'"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 6,
-   "id": "fecc9a00",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "image_size = 256\n",
-    "def image_transform(image):\n",
-    "    s = min(image.size)\n",
-    "    r = image_size / s\n",
-    "    s = (round(r * image.size[1]), round(r * image.size[0]))\n",
-    "    image = TF.resize(image, s, interpolation=InterpolationMode.LANCZOS)\n",
-    "    image = TF.center_crop(image, output_size = 2 * [image_size])\n",
-    "    image = torch.unsqueeze(T.ToTensor()(image), 0)\n",
-    "    image = image.permute(0, 2, 3, 1).numpy()\n",
-    "    return image"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 7,
-   "id": "4ce2211f",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "dataset = CaptionDataset(\n",
-    "    images_root=cc12m_images,\n",
-    "    captions_path=cc12m_list,\n",
-    "    image_transform=image_transform,\n",
-    "    image_transform_type='torchvision',\n",
-    "    include_captions=False\n",
-    ")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 8,
-   "id": "cc922704",
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "8592141"
-      ]
-     },
-     "execution_count": 8,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "len(dataset)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "62ad01c3",
-   "metadata": {},
-   "source": [
-    "## Encoding"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 9,
-   "id": "88f36d0b",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def encode(model, batch):\n",
-    "#     print(\"jitting encode function\")\n",
-    "    _, indices = model.encode(batch)\n",
-    "    return indices"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 10,
-   "id": "1f35f0cb",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def superbatch_generator(dataloader, num_tpus):\n",
-    "    iter_loader = iter(dataloader)\n",
-    "    for batch in iter_loader:\n",
-    "        superbatch = [batch.squeeze(1)]\n",
-    "        try:\n",
-    "            for b in range(num_tpus-1):\n",
-    "                batch = next(iter_loader)\n",
-    "                if batch is None:\n",
-    "                    break\n",
-    "                # Skip incomplete last batch\n",
-    "                if batch.shape[0] == dataloader.batch_size:\n",
-    "                    superbatch.append(batch.squeeze(1))\n",
-    "        except StopIteration:\n",
-    "            pass\n",
-    "        superbatch = torch.stack(superbatch, axis=0)\n",
-    "        yield superbatch"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 11,
-   "id": "2210705b",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import os\n",
-    "\n",
-    "def encode_captioned_dataset(dataset, output_tsv, batch_size=32, num_workers=16):\n",
-    "    if os.path.isfile(output_tsv):\n",
-    "        print(f\"Destination file {output_tsv} already exists, please move away.\")\n",
-    "        return\n",
-    "    \n",
-    "    num_tpus = 8    \n",
-    "    dataloader = DataLoader(dataset, batch_size=batch_size, num_workers=num_workers)\n",
-    "    superbatches = superbatch_generator(dataloader, num_tpus=num_tpus)\n",
-    "    \n",
-    "    p_encoder = pmap(lambda batch: encode(model, batch))\n",
-    "\n",
-    "    # We save each superbatch to avoid reallocation of buffers as we process them.\n",
-    "    # We keep the file open to prevent excessive file seeks.\n",
-    "    with open(output_tsv, \"w\") as file:\n",
-    "        iterations = len(dataset) // (batch_size * num_tpus)\n",
-    "        for n in tqdm(range(iterations)):\n",
-    "            superbatch = next(superbatches)\n",
-    "            encoded = p_encoder(superbatch.numpy())\n",
-    "            encoded = encoded.reshape(-1, encoded.shape[-1])\n",
-    "\n",
-    "            # Extract fields from the dataset internal `captions` property, and save to disk\n",
-    "            start_index = n * batch_size * num_tpus\n",
-    "            end_index = (n+1) * batch_size * num_tpus\n",
-    "            paths = dataset.captions[\"image_file\"][start_index:end_index].values\n",
-    "            captions = dataset.captions[\"caption\"][start_index:end_index].values\n",
-    "            encoded_as_string = list(map(lambda item: np.array2string(item, separator=',', max_line_width=50000, formatter={'int':lambda x: str(x)}), encoded))\n",
-    "            batch_df = pd.DataFrame.from_dict({\"image_file\": paths, \"caption\": captions, \"encoding\": encoded_as_string})\n",
-    "            batch_df.to_csv(file, sep='\\t', header=(n==0), index=None)\n",
-    "            "
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "7704863d",
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "  4%|██▋                                                                      | 621/16781 [07:09<3:02:46,  1.47it/s]"
-     ]
-    }
-   ],
-   "source": [
-    "encode_captioned_dataset(dataset, cc12m_output, batch_size=64, num_workers=16)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "8953dd84",
-   "metadata": {},
-   "source": [
-    "----"
-   ]
-  }
- ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.8.10"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 5
-}

dev/encoding/vqgan-jax-encoding-yfcc100m.ipynb

@@ -1,1129 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "markdown",
-   "id": "d0b72877",
-   "metadata": {},
-   "source": [
-    "# vqgan-jax-encoding-yfcc100m"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "ba7b31e6",
-   "metadata": {},
-   "source": [
-    "Same as `vqgan-jax-encoding-with-captions`, but for YFCC100M.\n",
-    "\n",
-    "This dataset was prepared by @borisdayma in Json lines format."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 92,
-   "id": "3b59489e",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import io\n",
-    "\n",
-    "import requests\n",
-    "from PIL import Image\n",
-    "import numpy as np\n",
-    "from tqdm import tqdm\n",
-    "\n",
-    "import torch\n",
-    "import torchvision.transforms as T\n",
-    "import torchvision.transforms.functional as TF\n",
-    "from torchvision.transforms import InterpolationMode\n",
-    "from torch.utils.data import Dataset, DataLoader\n",
-    "from torchvision.datasets.folder import default_loader\n",
-    "import os\n",
-    "\n",
-    "import jax\n",
-    "from jax import pmap"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "511c3b9e",
-   "metadata": {},
-   "source": [
-    "## VQGAN-JAX model"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 93,
-   "id": "2ca50dc7",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from vqgan_jax.modeling_flax_vqgan import VQModel"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "7b60da9a",
-   "metadata": {},
-   "source": [
-    "We'll use a VQGAN trained by using Taming Transformers and converted to a JAX model."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 167,
-   "id": "29ce8b15",
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Working with z of shape (1, 256, 16, 16) = 65536 dimensions.\n"
-     ]
-    }
-   ],
-   "source": [
-    "model = VQModel.from_pretrained(\"flax-community/vqgan_f16_16384\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "c7c4c1e6",
-   "metadata": {},
-   "source": [
-    "## Dataset"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 94,
-   "id": "33861477",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import pandas as pd\n",
-    "from pathlib import Path"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 134,
-   "id": "81b19eca",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "yfcc100m = Path('/home/khali/TPU-Test/YFCC100M_OpenAI_subset')\n",
-    "# Images are 'sharded' from the following directory\n",
-    "yfcc100m_images = yfcc100m/'data'/'data'/'images'\n",
-    "yfcc100m_metadata = yfcc100m/'metadata_YFCC100M.jsonl'\n",
-    "yfcc100m_output = yfcc100m/'metadata_encoded.tsv'"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "1c58bb4a",
-   "metadata": {},
-   "source": [
-    "### Cleanup"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "1a14ae3d",
-   "metadata": {},
-   "source": [
-    "We need to select entries with images that exist. Otherwise we can't build batches because `Dataloader` does not support `None` in batches. We use Huggingface Datasets, I understand they support threaded reading of jsonl files, and I was running out of memory when using pandas."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 96,
-   "id": "7811648c",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import datasets\n",
-    "from datasets import Dataset, load_dataset"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 10,
-   "id": "4811a230",
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "tcmalloc: large alloc 1254047744 bytes == 0xb2b08000 @  0x7f9e78632680 0x7f9e78653824 0x585b92 0x504d56 0x56acb6 0x5f5956 0x56acb6 0x5f5956 0x5a8cb3 0x56ae94 0x568d9a 0x68cdc7 0x5ff5d4 0x5c3cb0 0x56aadf 0x501148 0x56c422 0x501148 0x56c422 0x501148 0x504d56 0x56acb6 0x5f5956 0x56aadf 0x5f5956 0x56acb6 0x568d9a 0x5f5b33 0x50b7f8 0x5f2702 0x56c332\n",
-      "tcmalloc: large alloc 1254047744 bytes == 0xfd74e000 @  0x7f9e78632680 0x7f9e78653824 0x590214 0x586f90 0x56e1f3 0x5f5956 0x56acb6 0x5f5956 0x5a8cb3 0x56ae94 0x568d9a 0x68cdc7 0x5ff5d4 0x5c3cb0 0x56aadf 0x501148 0x56c422 0x501148 0x56c422 0x501148 0x504d56 0x56acb6 0x5f5956 0x56aadf 0x5f5956 0x56acb6 0x568d9a 0x5f5b33 0x50b7f8 0x5f2702 0x56c332\n",
-      "tcmalloc: large alloc 5016190976 bytes == 0x148b42000 @  0x7f9e78632680 0x7f9e78653824 0x5b9144 0x7f9b2929127e 0x7f9b29291a19 0x7f9b29291886 0x7f9b29291cef 0x7f9b2928f204 0x5f2cc9 0x5f30ff 0x5705f6 0x5f5956 0x56acb6 0x5f5956 0x56acb6 0x5f5956 0x56acb6 0x5f5956 0x5a8cb3 0x56ae94 0x568d9a 0x68cdc7 0x5ff5d4 0x5c3cb0 0x56aadf 0x501148 0x56c422 0x501148 0x56c422 0x501148 0x504d56\n",
-      "tcmalloc: large alloc 5019099136 bytes == 0x273f12000 @  0x7f9e78632680 0x7f9e78653824 0x5b9144 0x7f9b2929127e 0x7f9b29291a19 0x7f9b29291886 0x7f9b29291cef 0x7f9b2928f204 0x5f2cc9 0x5f30ff 0x5705f6 0x5f5956 0x56acb6 0x5f5956 0x56acb6 0x5f5956 0x56acb6 0x5f5956 0x5a8cb3 0x56ae94 0x568d9a 0x68cdc7 0x5ff5d4 0x5c3cb0 0x56aadf 0x501148 0x56c422 0x501148 0x56c422 0x501148 0x504d56\n",
-      "tcmalloc: large alloc 5019811840 bytes == 0x39f9a8000 @  0x7f9e78632680 0x7f9e78653824 0x5b9144 0x7f9b2929127e 0x7f9b29291a19 0x7f9b29291886 0x7f9b29291cef 0x7f9b2928f204 0x5f2cc9 0x5f30ff 0x5705f6 0x5f5956 0x56acb6 0x5f5956 0x56acb6 0x5f5956 0x56acb6 0x5f5956 0x5a8cb3 0x56ae94 0x568d9a 0x68cdc7 0x5ff5d4 0x5c3cb0 0x56aadf 0x501148 0x56c422 0x501148 0x56c422 0x501148 0x504d56\n",
-      "tcmalloc: large alloc 5024571392 bytes == 0x4cb4ec000 @  0x7f9e78632680 0x7f9e78653824 0x5b9144 0x7f9b2929127e 0x7f9b29291a19 0x7f9b29291886 0x7f9b29291cef 0x7f9b2928f204 0x5f2cc9 0x5f30ff 0x5705f6 0x5f5956 0x56acb6 0x5f5956 0x56acb6 0x5f5956 0x56acb6 0x5f5956 0x5a8cb3 0x56ae94 0x568d9a 0x68cdc7 0x5ff5d4 0x5c3cb0 0x56aadf 0x501148 0x56c422 0x501148 0x56c422 0x501148 0x504d56\n",
-      "tcmalloc: large alloc 5021097984 bytes == 0x4cb4ec000 @  0x7f9e78632680 0x7f9e78653824 0x5b9144 0x7f9b2929127e 0x7f9b29291a19 0x7f9b29291886 0x7f9b29291cef 0x7f9b2928f204 0x5f2cc9 0x5f30ff 0x5705f6 0x5f5956 0x56acb6 0x5f5956 0x56acb6 0x5f5956 0x56acb6 0x5f5956 0x5a8cb3 0x56ae94 0x568d9a 0x68cdc7 0x5ff5d4 0x5c3cb0 0x56aadf 0x501148 0x56c422 0x501148 0x56c422 0x501148 0x504d56\n",
-      "tcmalloc: large alloc 5022818304 bytes == 0x4cb4ec000 @  0x7f9e78632680 0x7f9e78653824 0x5b9144 0x7f9b2929127e 0x7f9b29291a19 0x7f9b29291886 0x7f9b29291cef 0x7f9b2928f204 0x5f2cc9 0x5f30ff 0x5705f6 0x5f5956 0x56acb6 0x5f5956 0x56acb6 0x5f5956 0x56acb6 0x5f5956 0x5a8cb3 0x56ae94 0x568d9a 0x68cdc7 0x5ff5d4 0x5c3cb0 0x56aadf 0x501148 0x56c422 0x501148 0x56c422 0x501148 0x504d56\n",
-      "tcmalloc: large alloc 5020794880 bytes == 0x4cb4ec000 @  0x7f9e78632680 0x7f9e78653824 0x5b9144 0x7f9b2929127e 0x7f9b29291a19 0x7f9b29291886 0x7f9b29291cef 0x7f9b2928f204 0x5f2cc9 0x5f30ff 0x5705f6 0x5f5956 0x56acb6 0x5f5956 0x56acb6 0x5f5956 0x56acb6 0x5f5956 0x5a8cb3 0x56ae94 0x568d9a 0x68cdc7 0x5ff5d4 0x5c3cb0 0x56aadf 0x501148 0x56c422 0x501148 0x56c422 0x501148 0x504d56\n",
-      "tcmalloc: large alloc 5019451392 bytes == 0x39f9a8000 @  0x7f9e78632680 0x7f9e78653824 0x5b9144 0x7f9b2929127e 0x7f9b29291a19 0x7f9b29291886 0x7f9b29291cef 0x7f9b2928f204 0x5f2cc9 0x5f30ff 0x5705f6 0x5f5956 0x56acb6 0x5f5956 0x56acb6 0x5f5956 0x56acb6 0x5f5956 0x5a8cb3 0x56ae94 0x568d9a 0x68cdc7 0x5ff5d4 0x5c3cb0 0x56aadf 0x501148 0x56c422 0x501148 0x56c422 0x501148 0x504d56\n",
-      "tcmalloc: large alloc 5020565504 bytes == 0x4cb4ec000 @  0x7f9e78632680 0x7f9e78653824 0x5b9144 0x7f9b2929127e 0x7f9b29291a19 0x7f9b29291886 0x7f9b29291cef 0x7f9b2928f204 0x5f2cc9 0x5f30ff 0x5705f6 0x5f5956 0x56acb6 0x5f5956 0x56acb6 0x5f5956 0x56acb6 0x5f5956 0x5a8cb3 0x56ae94 0x568d9a 0x68cdc7 0x5ff5d4 0x5c3cb0 0x56aadf 0x501148 0x56c422 0x501148 0x56c422 0x501148 0x504d56\n",
-      "tcmalloc: large alloc 5012561920 bytes == 0x273f12000 @  0x7f9e78632680 0x7f9e78653824 0x5b9144 0x7f9b2929127e 0x7f9b29291a19 0x7f9b29291886 0x7f9b29291cef 0x7f9b2928f204 0x5f2cc9 0x5f30ff 0x5705f6 0x5f5956 0x56acb6 0x5f5956 0x56acb6 0x5f5956 0x56acb6 0x5f5956 0x5a8cb3 0x56ae94 0x568d9a 0x68cdc7 0x5ff5d4 0x5c3cb0 0x56aadf 0x501148 0x56c422 0x501148 0x56c422 0x501148 0x504d56\n",
-      "tcmalloc: large alloc 5021835264 bytes == 0x5f6cba000 @  0x7f9e78632680 0x7f9e78653824 0x5b9144 0x7f9b2929127e 0x7f9b29291a19 0x7f9b29291886 0x7f9b29291cef 0x7f9b2928f204 0x5f2cc9 0x5f30ff 0x5705f6 0x5f5956 0x56acb6 0x5f5956 0x56acb6 0x5f5956 0x56acb6 0x5f5956 0x5a8cb3 0x56ae94 0x568d9a 0x68cdc7 0x5ff5d4 0x5c3cb0 0x56aadf 0x501148 0x56c422 0x501148 0x56c422 0x501148 0x504d56\n",
-      "tcmalloc: large alloc 5017436160 bytes == 0x273f12000 @  0x7f9e78632680 0x7f9e78653824 0x5b9144 0x7f9b2929127e 0x7f9b29291a19 0x7f9b29291886 0x7f9b29291cef 0x7f9b2928f204 0x5f2cc9 0x5f30ff 0x5705f6 0x5f5956 0x56acb6 0x5f5956 0x56acb6 0x5f5956 0x56acb6 0x5f5956 0x5a8cb3 0x56ae94 0x568d9a 0x68cdc7 0x5ff5d4 0x5c3cb0 0x56aadf 0x501148 0x56c422 0x501148 0x56c422 0x501148 0x504d56\n"
-     ]
-    }
-   ],
-   "source": [
-    "# The metadata is too bog to load into memory at once, so chopping it into chunks\n",
-    "chunk_size=1000000\n",
-    "batch_no=1\n",
-    "for chunk in pd.read_json(yfcc100m_metadata, orient=\"records\", lines=True,chunksize=chunk_size):\n",
-    "    chunk.to_csv('./chunks/chunk'+str(batch_no)+'.tsv', sep=\"\\t\", index=False)\n",
-    "    batch_no+=1"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 25,
-   "id": "46b2f083",
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/html": [
-       "<div>\n",
-       "<style scoped>\n",
-       "    .dataframe tbody tr th:only-of-type {\n",
-       "        vertical-align: middle;\n",
-       "    }\n",
-       "\n",
-       "    .dataframe tbody tr th {\n",
-       "        vertical-align: top;\n",
-       "    }\n",
-       "\n",
-       "    .dataframe thead th {\n",
-       "        text-align: right;\n",
-       "    }\n",
-       "</style>\n",
-       "<table border=\"1\" class=\"dataframe\">\n",
-       "  <thead>\n",
-       "    <tr style=\"text-align: right;\">\n",
-       "      <th></th>\n",
-       "      <th>photoid</th>\n",
-       "      <th>uid</th>\n",
-       "      <th>unickname</th>\n",
-       "      <th>datetaken</th>\n",
-       "      <th>dateuploaded</th>\n",
-       "      <th>capturedevice</th>\n",
-       "      <th>title</th>\n",
-       "      <th>description</th>\n",
-       "      <th>usertags</th>\n",
-       "      <th>machinetags</th>\n",
-       "      <th>...</th>\n",
-       "      <th>licenseurl</th>\n",
-       "      <th>serverid</th>\n",
-       "      <th>farmid</th>\n",
-       "      <th>secret</th>\n",
-       "      <th>secretoriginal</th>\n",
-       "      <th>ext</th>\n",
-       "      <th>marker</th>\n",
-       "      <th>key</th>\n",
-       "      <th>title_clean</th>\n",
-       "      <th>description_clean</th>\n",
-       "    </tr>\n",
-       "  </thead>\n",
-       "  <tbody>\n",
-       "    <tr>\n",
-       "      <th>0</th>\n",
-       "      <td>137943</td>\n",
-       "      <td>48600072071@N01</td>\n",
-       "      <td>doctor+paradox</td>\n",
-       "      <td>2004-08-01 18:13:06.0</td>\n",
-       "      <td>1091409186</td>\n",
-       "      <td>NaN</td>\n",
-       "      <td>A+Picture+Share%21</td>\n",
-       "      <td>Antenna</td>\n",
-       "      <td>cameraphone,cayugaheights,green,hydrant,ithaca...</td>\n",
-       "      <td>NaN</td>\n",
-       "      <td>...</td>\n",
-       "      <td>http://creativecommons.org/licenses/by-nc-sa/2.0/</td>\n",
-       "      <td>1</td>\n",
-       "      <td>1</td>\n",
-       "      <td>1650c7cdc6</td>\n",
-       "      <td>1650c7cdc6</td>\n",
-       "      <td>jpg</td>\n",
-       "      <td>0</td>\n",
-       "      <td>d29e7c6a3028418c64eb15e3cf577c2</td>\n",
-       "      <td>A Picture Share!</td>\n",
-       "      <td>Antenna</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>1</th>\n",
-       "      <td>1246361</td>\n",
-       "      <td>44124324682@N01</td>\n",
-       "      <td>mharrsch</td>\n",
-       "      <td>2004-11-03 23:04:02.0</td>\n",
-       "      <td>1099523042</td>\n",
-       "      <td>NaN</td>\n",
-       "      <td>An+ornate+Roman+urn</td>\n",
-       "      <td>Photographed+at+the+%3Ca+href%3D%22http%3A%2F%...</td>\n",
-       "      <td>ancient,baltimore,burial,death,empire,funeral,...</td>\n",
-       "      <td>NaN</td>\n",
-       "      <td>...</td>\n",
-       "      <td>http://creativecommons.org/licenses/by-nc-sa/2.0/</td>\n",
-       "      <td>1</td>\n",
-       "      <td>1</td>\n",
-       "      <td>cf37054610</td>\n",
-       "      <td>cf37054610</td>\n",
-       "      <td>jpg</td>\n",
-       "      <td>0</td>\n",
-       "      <td>d29f01b149167d683f9ddde464bb3db</td>\n",
-       "      <td>An ornate Roman urn</td>\n",
-       "      <td>Photographed at the Walters Art Museum, Baltim...</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>2</th>\n",
-       "      <td>1251599</td>\n",
-       "      <td>51035803024@N01</td>\n",
-       "      <td>bmitd67</td>\n",
-       "      <td>2004-10-30 17:09:32.0</td>\n",
-       "      <td>1099538888</td>\n",
-       "      <td>Canon+PowerShot+S30</td>\n",
-       "      <td>Jai+%26+Tara+on+the+Cumberland</td>\n",
-       "      <td>Another+trip+for+the+happy+couple.</td>\n",
-       "      <td>blue+heron,cumberland+river,jai,tara,tennessee</td>\n",
-       "      <td>NaN</td>\n",
-       "      <td>...</td>\n",
-       "      <td>http://creativecommons.org/licenses/by-nc-sa/2.0/</td>\n",
-       "      <td>1</td>\n",
-       "      <td>1</td>\n",
-       "      <td>4a4234e32c</td>\n",
-       "      <td>4a4234e32c</td>\n",
-       "      <td>jpg</td>\n",
-       "      <td>0</td>\n",
-       "      <td>d296e9e34bdae41edb6c679ff824ab2a</td>\n",
-       "      <td>Jai &amp; Tara on the Cumberland</td>\n",
-       "      <td>Another trip for the happy couple.</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>3</th>\n",
-       "      <td>2348587</td>\n",
-       "      <td>73621375@N00</td>\n",
-       "      <td>Thom+Watson</td>\n",
-       "      <td>2004-12-18 21:08:09.0</td>\n",
-       "      <td>1103497228</td>\n",
-       "      <td>SONY+DSC-W1</td>\n",
-       "      <td>Castle+gate+-+%22lite-brited%22</td>\n",
-       "      <td>Taken+at+the+Miracle+of+Lights+display+in+Cent...</td>\n",
-       "      <td>bullrunpark,castle,centreville,christmas,decor...</td>\n",
-       "      <td>NaN</td>\n",
-       "      <td>...</td>\n",
-       "      <td>http://creativecommons.org/licenses/by-nc-sa/2.0/</td>\n",
-       "      <td>2</td>\n",
-       "      <td>1</td>\n",
-       "      <td>7162c974c3</td>\n",
-       "      <td>7162c974c3</td>\n",
-       "      <td>jpg</td>\n",
-       "      <td>0</td>\n",
-       "      <td>d29ce96395848478b1e8396e44899</td>\n",
-       "      <td>Castle gate - \"lite-brited\"</td>\n",
-       "      <td>Taken at the Miracle of Lights display in Cent...</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>4</th>\n",
-       "      <td>3516047</td>\n",
-       "      <td>48600072071@N01</td>\n",
-       "      <td>doctor+paradox</td>\n",
-       "      <td>2005-01-18 16:44:18.0</td>\n",
-       "      <td>1106084658</td>\n",
-       "      <td>NaN</td>\n",
-       "      <td>A+Picture+Share%21</td>\n",
-       "      <td>Tabular</td>\n",
-       "      <td>cameraphone,moblog,unfound</td>\n",
-       "      <td>NaN</td>\n",
-       "      <td>...</td>\n",
-       "      <td>http://creativecommons.org/licenses/by-nc-sa/2.0/</td>\n",
-       "      <td>3</td>\n",
-       "      <td>1</td>\n",
-       "      <td>663e0d8b3d</td>\n",
-       "      <td>663e0d8b3d</td>\n",
-       "      <td>jpg</td>\n",
-       "      <td>0</td>\n",
-       "      <td>d29abf32c4e12ff881f975b70e0cec0</td>\n",
-       "      <td>A Picture Share!</td>\n",
-       "      <td>Tabular</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>...</th>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>999995</th>\n",
-       "      <td>4648651054</td>\n",
-       "      <td>24511045@N04</td>\n",
-       "      <td>mtfrazier</td>\n",
-       "      <td>2010-05-02 15:47:45.0</td>\n",
-       "      <td>1275083371</td>\n",
-       "      <td>Canon+EOS+50D</td>\n",
-       "      <td>U.S.+Navy+Blue+Angels%3A+2010</td>\n",
-       "      <td>2+May+2010%0ASunday%0ASt.+Joseph%2C+Missouri</td>\n",
-       "      <td>NaN</td>\n",
-       "      <td>NaN</td>\n",
-       "      <td>...</td>\n",
-       "      <td>http://creativecommons.org/licenses/by-nc-nd/2.0/</td>\n",
-       "      <td>4072</td>\n",
-       "      <td>5</td>\n",
-       "      <td>2d12d73fb0</td>\n",
-       "      <td>dd5856ea42</td>\n",
-       "      <td>jpg</td>\n",
-       "      <td>0</td>\n",
-       "      <td>60fa2911cb81eb25b356e9fee978aef</td>\n",
-       "      <td>U.S. Navy Blue Angels: 2010</td>\n",
-       "      <td>2 May 2010 Sunday St. Joseph, Missouri</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>999996</th>\n",
-       "      <td>4652130996</td>\n",
-       "      <td>21963865@N04</td>\n",
-       "      <td>GRAB1.0</td>\n",
-       "      <td>2010-05-29 19:23:10.0</td>\n",
-       "      <td>1275200833</td>\n",
-       "      <td>SONY+DSLR-A230</td>\n",
-       "      <td>Attempts+on+Her+Life</td>\n",
-       "      <td>BAPA+1+production+of+Martin+Crimp%27s+Attempts...</td>\n",
-       "      <td>NaN</td>\n",
-       "      <td>NaN</td>\n",
-       "      <td>...</td>\n",
-       "      <td>http://creativecommons.org/licenses/by-nc-nd/2.0/</td>\n",
-       "      <td>4003</td>\n",
-       "      <td>5</td>\n",
-       "      <td>8889121579</td>\n",
-       "      <td>2f46599456</td>\n",
-       "      <td>jpg</td>\n",
-       "      <td>0</td>\n",
-       "      <td>60f5ef5ce4c2d24566226abebd67d4</td>\n",
-       "      <td>Attempts on Her Life</td>\n",
-       "      <td>BAPA 1 production of Martin Crimp's Attempts o...</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>999997</th>\n",
-       "      <td>4652568339</td>\n",
-       "      <td>64025277@N00</td>\n",
-       "      <td>1Sock</td>\n",
-       "      <td>2010-05-13 15:38:37.0</td>\n",
-       "      <td>1275234267</td>\n",
-       "      <td>Canon+EOS+DIGITAL+REBEL+XT</td>\n",
-       "      <td>Carlsbad+Caverns+3</td>\n",
-       "      <td>%E2%99%A5%E2%99%A5%E2%99%A5%E2%99%A5%E2%99%A5%...</td>\n",
-       "      <td>carlsbad,carlsbad+caverns,cave,faa,new+mexico,...</td>\n",
-       "      <td>NaN</td>\n",
-       "      <td>...</td>\n",
-       "      <td>http://creativecommons.org/licenses/by-nc-nd/2.0/</td>\n",
-       "      <td>4010</td>\n",
-       "      <td>5</td>\n",
-       "      <td>0a1808a69e</td>\n",
-       "      <td>cf6d348e3d</td>\n",
-       "      <td>jpg</td>\n",
-       "      <td>0</td>\n",
-       "      <td>60f029482d1d1028fda5281daf498f</td>\n",
-       "      <td>Carlsbad Caverns 3</td>\n",
-       "      <td>♥♥♥♥♥♥♥ Interested in purchasing this photogra...</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>999998</th>\n",
-       "      <td>4653110895</td>\n",
-       "      <td>20483509@N00</td>\n",
-       "      <td>subberculture</td>\n",
-       "      <td>2010-05-30 15:37:05.0</td>\n",
-       "      <td>1275245596</td>\n",
-       "      <td>Canon+DIGITAL+IXUS+40</td>\n",
-       "      <td>Want</td>\n",
-       "      <td>Isn%27t+that+gorgeous%3F</td>\n",
-       "      <td>2010,edinburgh+museum,may,phonebox,wood</td>\n",
-       "      <td>NaN</td>\n",
-       "      <td>...</td>\n",
-       "      <td>http://creativecommons.org/licenses/by-sa/2.0/</td>\n",
-       "      <td>4066</td>\n",
-       "      <td>5</td>\n",
-       "      <td>77c3b3a254</td>\n",
-       "      <td>c4697e1511</td>\n",
-       "      <td>jpg</td>\n",
-       "      <td>0</td>\n",
-       "      <td>60f72775f433cf8de3efaeb431866153</td>\n",
-       "      <td>Want</td>\n",
-       "      <td>Isn't that gorgeous?</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>999999</th>\n",
-       "      <td>4655503987</td>\n",
-       "      <td>8457193@N07</td>\n",
-       "      <td>zackojones</td>\n",
-       "      <td>2010-05-30 15:34:58.0</td>\n",
-       "      <td>1275310230</td>\n",
-       "      <td>Canon+EOS+7D</td>\n",
-       "      <td>Summertime</td>\n",
-       "      <td>You+gotta+love+it%21</td>\n",
-       "      <td>georgia,savannah,united+states,us</td>\n",
-       "      <td>NaN</td>\n",
-       "      <td>...</td>\n",
-       "      <td>http://creativecommons.org/licenses/by-nc-sa/2.0/</td>\n",
-       "      <td>4043</td>\n",
-       "      <td>5</td>\n",
-       "      <td>caff543bfe</td>\n",
-       "      <td>f60952ac4d</td>\n",
-       "      <td>jpg</td>\n",
-       "      <td>0</td>\n",
-       "      <td>60f687e11b913bce461e9525d8047e0</td>\n",
-       "      <td>Summertime</td>\n",
-       "      <td>You gotta love it!</td>\n",
-       "    </tr>\n",
-       "  </tbody>\n",
-       "</table>\n",
-       "<p>1000000 rows × 26 columns</p>\n",
-       "</div>"
-      ],
-      "text/plain": [
-       "           photoid              uid       unickname              datetaken  \\\n",
-       "0           137943  48600072071@N01  doctor+paradox  2004-08-01 18:13:06.0   \n",
-       "1          1246361  44124324682@N01        mharrsch  2004-11-03 23:04:02.0   \n",
-       "2          1251599  51035803024@N01         bmitd67  2004-10-30 17:09:32.0   \n",
-       "3          2348587     73621375@N00     Thom+Watson  2004-12-18 21:08:09.0   \n",
-       "4          3516047  48600072071@N01  doctor+paradox  2005-01-18 16:44:18.0   \n",
-       "...            ...              ...             ...                    ...   \n",
-       "999995  4648651054     24511045@N04       mtfrazier  2010-05-02 15:47:45.0   \n",
-       "999996  4652130996     21963865@N04         GRAB1.0  2010-05-29 19:23:10.0   \n",
-       "999997  4652568339     64025277@N00           1Sock  2010-05-13 15:38:37.0   \n",
-       "999998  4653110895     20483509@N00   subberculture  2010-05-30 15:37:05.0   \n",
-       "999999  4655503987      8457193@N07      zackojones  2010-05-30 15:34:58.0   \n",
-       "\n",
-       "        dateuploaded               capturedevice  \\\n",
-       "0         1091409186                         NaN   \n",
-       "1         1099523042                         NaN   \n",
-       "2         1099538888         Canon+PowerShot+S30   \n",
-       "3         1103497228                 SONY+DSC-W1   \n",
-       "4         1106084658                         NaN   \n",
-       "...              ...                         ...   \n",
-       "999995    1275083371               Canon+EOS+50D   \n",
-       "999996    1275200833              SONY+DSLR-A230   \n",
-       "999997    1275234267  Canon+EOS+DIGITAL+REBEL+XT   \n",
-       "999998    1275245596       Canon+DIGITAL+IXUS+40   \n",
-       "999999    1275310230                Canon+EOS+7D   \n",
-       "\n",
-       "                                  title  \\\n",
-       "0                    A+Picture+Share%21   \n",
-       "1                   An+ornate+Roman+urn   \n",
-       "2        Jai+%26+Tara+on+the+Cumberland   \n",
-       "3       Castle+gate+-+%22lite-brited%22   \n",
-       "4                    A+Picture+Share%21   \n",
-       "...                                 ...   \n",
-       "999995    U.S.+Navy+Blue+Angels%3A+2010   \n",
-       "999996             Attempts+on+Her+Life   \n",
-       "999997               Carlsbad+Caverns+3   \n",
-       "999998                             Want   \n",
-       "999999                       Summertime   \n",
-       "\n",
-       "                                              description  \\\n",
-       "0                                                 Antenna   \n",
-       "1       Photographed+at+the+%3Ca+href%3D%22http%3A%2F%...   \n",
-       "2                      Another+trip+for+the+happy+couple.   \n",
-       "3       Taken+at+the+Miracle+of+Lights+display+in+Cent...   \n",
-       "4                                                 Tabular   \n",
-       "...                                                   ...   \n",
-       "999995       2+May+2010%0ASunday%0ASt.+Joseph%2C+Missouri   \n",
-       "999996  BAPA+1+production+of+Martin+Crimp%27s+Attempts...   \n",
-       "999997  %E2%99%A5%E2%99%A5%E2%99%A5%E2%99%A5%E2%99%A5%...   \n",
-       "999998                           Isn%27t+that+gorgeous%3F   \n",
-       "999999                               You+gotta+love+it%21   \n",
-       "\n",
-       "                                                 usertags machinetags  ...  \\\n",
-       "0       cameraphone,cayugaheights,green,hydrant,ithaca...         NaN  ...   \n",
-       "1       ancient,baltimore,burial,death,empire,funeral,...         NaN  ...   \n",
-       "2          blue+heron,cumberland+river,jai,tara,tennessee         NaN  ...   \n",
-       "3       bullrunpark,castle,centreville,christmas,decor...         NaN  ...   \n",
-       "4                              cameraphone,moblog,unfound         NaN  ...   \n",
-       "...                                                   ...         ...  ...   \n",
-       "999995                                                NaN         NaN  ...   \n",
-       "999996                                                NaN         NaN  ...   \n",
-       "999997  carlsbad,carlsbad+caverns,cave,faa,new+mexico,...         NaN  ...   \n",
-       "999998            2010,edinburgh+museum,may,phonebox,wood         NaN  ...   \n",
-       "999999                  georgia,savannah,united+states,us         NaN  ...   \n",
-       "\n",
-       "                                               licenseurl  serverid  farmid  \\\n",
-       "0       http://creativecommons.org/licenses/by-nc-sa/2.0/         1       1   \n",
-       "1       http://creativecommons.org/licenses/by-nc-sa/2.0/         1       1   \n",
-       "2       http://creativecommons.org/licenses/by-nc-sa/2.0/         1       1   \n",
-       "3       http://creativecommons.org/licenses/by-nc-sa/2.0/         2       1   \n",
-       "4       http://creativecommons.org/licenses/by-nc-sa/2.0/         3       1   \n",
-       "...                                                   ...       ...     ...   \n",
-       "999995  http://creativecommons.org/licenses/by-nc-nd/2.0/      4072       5   \n",
-       "999996  http://creativecommons.org/licenses/by-nc-nd/2.0/      4003       5   \n",
-       "999997  http://creativecommons.org/licenses/by-nc-nd/2.0/      4010       5   \n",
-       "999998     http://creativecommons.org/licenses/by-sa/2.0/      4066       5   \n",
-       "999999  http://creativecommons.org/licenses/by-nc-sa/2.0/      4043       5   \n",
-       "\n",
-       "            secret secretoriginal  ext marker  \\\n",
-       "0       1650c7cdc6     1650c7cdc6  jpg      0   \n",
-       "1       cf37054610     cf37054610  jpg      0   \n",
-       "2       4a4234e32c     4a4234e32c  jpg      0   \n",
-       "3       7162c974c3     7162c974c3  jpg      0   \n",
-       "4       663e0d8b3d     663e0d8b3d  jpg      0   \n",
-       "...            ...            ...  ...    ...   \n",
-       "999995  2d12d73fb0     dd5856ea42  jpg      0   \n",
-       "999996  8889121579     2f46599456  jpg      0   \n",
-       "999997  0a1808a69e     cf6d348e3d  jpg      0   \n",
-       "999998  77c3b3a254     c4697e1511  jpg      0   \n",
-       "999999  caff543bfe     f60952ac4d  jpg      0   \n",
-       "\n",
-       "                                     key                   title_clean  \\\n",
-       "0        d29e7c6a3028418c64eb15e3cf577c2              A Picture Share!   \n",
-       "1        d29f01b149167d683f9ddde464bb3db           An ornate Roman urn   \n",
-       "2       d296e9e34bdae41edb6c679ff824ab2a  Jai & Tara on the Cumberland   \n",
-       "3          d29ce96395848478b1e8396e44899   Castle gate - \"lite-brited\"   \n",
-       "4        d29abf32c4e12ff881f975b70e0cec0              A Picture Share!   \n",
-       "...                                  ...                           ...   \n",
-       "999995   60fa2911cb81eb25b356e9fee978aef   U.S. Navy Blue Angels: 2010   \n",
-       "999996    60f5ef5ce4c2d24566226abebd67d4          Attempts on Her Life   \n",
-       "999997    60f029482d1d1028fda5281daf498f            Carlsbad Caverns 3   \n",
-       "999998  60f72775f433cf8de3efaeb431866153                          Want   \n",
-       "999999   60f687e11b913bce461e9525d8047e0                    Summertime   \n",
-       "\n",
-       "                                        description_clean  \n",
-       "0                                                 Antenna  \n",
-       "1       Photographed at the Walters Art Museum, Baltim...  \n",
-       "2                      Another trip for the happy couple.  \n",
-       "3       Taken at the Miracle of Lights display in Cent...  \n",
-       "4                                                 Tabular  \n",
-       "...                                                   ...  \n",
-       "999995             2 May 2010 Sunday St. Joseph, Missouri  \n",
-       "999996  BAPA 1 production of Martin Crimp's Attempts o...  \n",
-       "999997  ♥♥♥♥♥♥♥ Interested in purchasing this photogra...  \n",
-       "999998                               Isn't that gorgeous?  \n",
-       "999999                                 You gotta love it!  \n",
-       "\n",
-       "[1000000 rows x 26 columns]"
-      ]
-     },
-     "execution_count": 25,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "# looking up at a chunk\n",
-    "pd.read_csv(\"./chunks/chunk1.tsv\", sep=\"\\t\")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 98,
-   "id": "c51c5597",
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/html": [
-       "<div>\n",
-       "<style scoped>\n",
-       "    .dataframe tbody tr th:only-of-type {\n",
-       "        vertical-align: middle;\n",
-       "    }\n",
-       "\n",
-       "    .dataframe tbody tr th {\n",
-       "        vertical-align: top;\n",
-       "    }\n",
-       "\n",
-       "    .dataframe thead th {\n",
-       "        text-align: right;\n",
-       "    }\n",
-       "</style>\n",
-       "<table border=\"1\" class=\"dataframe\">\n",
-       "  <thead>\n",
-       "    <tr style=\"text-align: right;\">\n",
-       "      <th></th>\n",
-       "      <th>key</th>\n",
-       "      <th>title_clean</th>\n",
-       "      <th>description_clean</th>\n",
-       "      <th>ext</th>\n",
-       "    </tr>\n",
-       "  </thead>\n",
-       "  <tbody>\n",
-       "    <tr>\n",
-       "      <th>0</th>\n",
-       "      <td>d29e7c6a3028418c64eb15e3cf577c2</td>\n",
-       "      <td>A Picture Share!</td>\n",
-       "      <td>Antenna</td>\n",
-       "      <td>jpg</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>1</th>\n",
-       "      <td>d29f01b149167d683f9ddde464bb3db</td>\n",
-       "      <td>An ornate Roman urn</td>\n",
-       "      <td>Photographed at the Walters Art Museum, Baltim...</td>\n",
-       "      <td>jpg</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>2</th>\n",
-       "      <td>d296e9e34bdae41edb6c679ff824ab2a</td>\n",
-       "      <td>Jai &amp; Tara on the Cumberland</td>\n",
-       "      <td>Another trip for the happy couple.</td>\n",
-       "      <td>jpg</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>3</th>\n",
-       "      <td>d29ce96395848478b1e8396e44899</td>\n",
-       "      <td>Castle gate - \"lite-brited\"</td>\n",
-       "      <td>Taken at the Miracle of Lights display in Cent...</td>\n",
-       "      <td>jpg</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>4</th>\n",
-       "      <td>d29abf32c4e12ff881f975b70e0cec0</td>\n",
-       "      <td>A Picture Share!</td>\n",
-       "      <td>Tabular</td>\n",
-       "      <td>jpg</td>\n",
-       "    </tr>\n",
-       "  </tbody>\n",
-       "</table>\n",
-       "</div>"
-      ],
-      "text/plain": [
-       "                                key                   title_clean  \\\n",
-       "0   d29e7c6a3028418c64eb15e3cf577c2              A Picture Share!   \n",
-       "1   d29f01b149167d683f9ddde464bb3db           An ornate Roman urn   \n",
-       "2  d296e9e34bdae41edb6c679ff824ab2a  Jai & Tara on the Cumberland   \n",
-       "3     d29ce96395848478b1e8396e44899   Castle gate - \"lite-brited\"   \n",
-       "4   d29abf32c4e12ff881f975b70e0cec0              A Picture Share!   \n",
-       "\n",
-       "                                   description_clean  ext  \n",
-       "0                                            Antenna  jpg  \n",
-       "1  Photographed at the Walters Art Museum, Baltim...  jpg  \n",
-       "2                 Another trip for the happy couple.  jpg  \n",
-       "3  Taken at the Miracle of Lights display in Cent...  jpg  \n",
-       "4                                            Tabular  jpg  "
-      ]
-     },
-     "execution_count": 98,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "# Looking at a chunk with only the relevant columns that we need\n",
-    "df = pd.read_csv(\"./chunks/chunk1.tsv\", sep=\"\\t\")[[\"key\", \"title_clean\", \"description_clean\", \"ext\"]]\n",
-    "df.head()"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "cc1668f8",
-   "metadata": {},
-   "source": [
-    "### Grabbing each chunks from the folder, cleaning it up, only taking the entries which image exist and appending it to the global df"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "abbcccf3",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# the function that helps us to decide whether an image with certain id exists in storage, we only take the ones that we have the images for\n",
-    "def image_exists(item):\n",
-    "    name, _, _, ext, _ = item\n",
-    "    root=str(yfcc100m_images)\n",
-    "    image_path = (Path(root)/name[0:3]/name[3:6]/name).with_suffix(\".\"+ext)\n",
-    "    if image_path.exists():\n",
-    "        return True\n",
-    "    else:\n",
-    "        return None"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 86,
-   "id": "44fa86ab",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# This cell does it all, grabs each chunk, cleans it up based on image existing condition, etc.\n",
-    "global_df = pd.DataFrame()\n",
-    "chunks_dir = \"./chunks\"\n",
-    "for filename in os.listdir(chunks_dir):\n",
-    "        df = pd.read_csv(f\"./chunks/{str(filename)}\", sep=\"\\t\")[[\"key\", \"title_clean\", \"description_clean\", \"ext\"]]\n",
-    "        df['caption'] = df[\"title_clean\"]+\". \"+df['description_clean']\n",
-    "        df['is_exist'] = df.apply(image_exists, axis=1)\n",
-    "        df = df.dropna()[[\"key\", \"caption\"]]\n",
-    "        df.columns = ['image_file', 'caption']\n",
-    "        global_df = global_df.append(df, ignore_index=True)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 89,
-   "id": "45024fdc",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# saving the tsv to disk\n",
-    "global_df.to_csv('./chunks/YFCC_subset_clean.tsv', sep=\"\\t\", index=False)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 101,
-   "id": "dca4eb73",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# loading the tsv from disk (for explicitness, also my electricity was gone, glad it happened after I saved to the disk :( )\n",
-    "\n",
-    "dataset = pd.read_csv(f\"./chunks/YFCC_subset_clean.tsv\", sep=\"\\t\")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 153,
-   "id": "a511264a",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "\"\"\"\n",
-    "Luke Melas-Kyriazi's dataset.py's modified version for YFCC\n",
-    "\"\"\"\n",
-    "import warnings\n",
-    "from typing import Optional, Callable\n",
-    "from pathlib import Path\n",
-    "import numpy as np\n",
-    "import torch\n",
-    "import pandas as pd\n",
-    "from torch.utils.data import Dataset\n",
-    "from torchvision.datasets.folder import default_loader\n",
-    "from PIL import ImageFile\n",
-    "from PIL.Image import DecompressionBombWarning\n",
-    "ImageFile.LOAD_TRUNCATED_IMAGES = True\n",
-    "warnings.filterwarnings(\"ignore\", category=UserWarning)\n",
-    "warnings.filterwarnings(\"ignore\", category=DecompressionBombWarning)\n",
-    "\n",
-    "\n",
-    "class CaptionDataset(Dataset):\n",
-    "    \"\"\"\n",
-    "    A PyTorch Dataset class for (image, texts) tasks. Note that this dataset \n",
-    "    returns the raw text rather than tokens. This is done on purpose, because\n",
-    "    it's easy to tokenize a batch of text after loading it from this dataset.\n",
-    "    \"\"\"\n",
-    "\n",
-    "    def __init__(self, *, images_root: str, captions_path: str, text_transform: Optional[Callable] = None, \n",
-    "                 image_transform: Optional[Callable] = None, image_transform_type: str = 'torchvision',\n",
-    "                 include_captions: bool = True):\n",
-    "        \"\"\"\n",
-    "        :param images_root: folder where images are stored\n",
-    "        :param captions_path: path to csv that maps image filenames to captions\n",
-    "        :param image_transform: image transform pipeline\n",
-    "        :param text_transform: image transform pipeline\n",
-    "        :param image_transform_type: image transform type, either `torchvision` or `albumentations`\n",
-    "        :param include_captions: Returns a dictionary with `image`, `text` if `true`; otherwise returns just the images.\n",
-    "        \"\"\"\n",
-    "\n",
-    "        # Base path for images\n",
-    "        self.images_root = Path(images_root)\n",
-    "\n",
-    "        # Load captions as DataFrame\n",
-    "        self.captions = pd.read_csv(f\"./chunks/YFCC_subset_clean.tsv\", sep=\"\\t\")\n",
-    "        self.captions['image_file'] = self.captions['image_file'].astype(str)\n",
-    "\n",
-    "        # PyTorch transformation pipeline for the image (normalizing, etc.)\n",
-    "        self.text_transform = text_transform\n",
-    "        self.image_transform = image_transform\n",
-    "        self.image_transform_type = image_transform_type.lower()\n",
-    "        assert self.image_transform_type in ['torchvision', 'albumentations']\n",
-    "\n",
-    "        # Total number of datapoints\n",
-    "        self.size = len(self.captions)\n",
-    "\n",
-    "        # Return image+captions or just images\n",
-    "        self.include_captions = include_captions\n",
-    "    \n",
-    "    def image_exists(item):\n",
-    "        name, caption = item\n",
-    "        root=str(self.images_root)\n",
-    "        image_path = (Path(root)/name[0:3]/name[3:6]/name).with_suffix(\".jpg\")\n",
-    "\n",
-    "        return image_path.exists()\n",
-    "\n",
-    "    def verify_that_all_images_exist(self):\n",
-    "        for image_file in self.captions['image_file']:\n",
-    "            if not image_exists:\n",
-    "                print(f'file does not exist: {p}')\n",
-    "\n",
-    "    def _get_raw_image(self, i):\n",
-    "        name = self.captions.iloc[i]['image_file']\n",
-    "        image_path = (Path(self.images_root)/name[0:3]/name[3:6]/name).with_suffix(\".jpg\")\n",
-    "        image = default_loader(image_path)\n",
-    "        return image\n",
-    "\n",
-    "    def _get_raw_text(self, i):\n",
-    "        return self.captions.iloc[i]['caption']\n",
-    "\n",
-    "    def __getitem__(self, i):\n",
-    "        image = self._get_raw_image(i)\n",
-    "        caption = self._get_raw_text(i)\n",
-    "        if self.image_transform is not None:\n",
-    "            if self.image_transform_type == 'torchvision':\n",
-    "                image = self.image_transform(image)\n",
-    "            elif self.image_transform_type == 'albumentations':\n",
-    "                image = self.image_transform(image=np.array(image))['image']\n",
-    "            else:\n",
-    "                raise NotImplementedError(f\"{self.image_transform_type=}\")\n",
-    "        return {'image': image, 'text': caption} if self.include_captions else image\n",
-    "\n",
-    "    def __len__(self):\n",
-    "        return self.size\n",
-    "\n",
-    "\n",
-    "if __name__ == \"__main__\":\n",
-    "    import albumentations as A\n",
-    "    from albumentations.pytorch import ToTensorV2\n",
-    "    from transformers import AutoTokenizer\n",
-    "    \n",
-    "\n",
-    "    images_root = \"/home/khali/TPU-Test/YFCC100M_OpenAI_subset/data/data/images\"\n",
-    "    captions_path = './YFCC_subset_clean.tsv'\n",
-    "    image_size = 256\n",
-    "    \n",
-    "    # Create transforms\n",
-    "    def image_transform(image):\n",
-    "        s = min(image.size)\n",
-    "        r = image_size / s\n",
-    "        s = (round(r * image.size[1]), round(r * image.size[0]))\n",
-    "        image = TF.resize(image, s, interpolation=InterpolationMode.LANCZOS)\n",
-    "        image = TF.center_crop(image, output_size = 2 * [image_size])\n",
-    "        image = torch.unsqueeze(T.ToTensor()(image), 0)\n",
-    "        image = image.permute(0, 2, 3, 1).numpy()\n",
-    "        return image\n",
-    "    \n",
-    "    # Create dataset\n",
-    "    dataset = CaptionDataset(\n",
-    "        images_root=images_root,\n",
-    "        captions_path=captions_path,\n",
-    "        image_transform=image_transform,\n",
-    "        image_transform_type='torchvision',\n",
-    "        include_captions=False\n",
-    "    )"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 155,
-   "id": "cc922704",
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "2483316"
-      ]
-     },
-     "execution_count": 155,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "len(dataset)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 156,
-   "id": "6e47ba46",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "dataloader = DataLoader(dataset, batch_size=32, num_workers=4)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 1,
-   "id": "c8a130eb",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# looking at a batch\n",
-    "next(iter(dataloader))"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "c192fd44",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# import matplotlib.pyplot as plt\n",
-    "# for tensor_image, _ in dataloader:\n",
-    "#     print(tensor_image)\n",
-    "#     plt.imshow(tensor_image.permute(1, 2, 0))\n",
-    "#     break"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "62ad01c3",
-   "metadata": {},
-   "source": [
-    "## Encoding"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 158,
-   "id": "88f36d0b",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def encode(model, batch):\n",
-    "#     print(\"jitting encode function\")\n",
-    "    _, indices = model.encode(batch)\n",
-    "    return indices"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 160,
-   "id": "1f35f0cb",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def superbatch_generator(dataloader, num_tpus):\n",
-    "    iter_loader = iter(dataloader)\n",
-    "    for batch in iter_loader:\n",
-    "        superbatch = [batch.squeeze(1)]\n",
-    "        try:\n",
-    "            for b in range(num_tpus-1):\n",
-    "                batch = next(iter_loader)\n",
-    "                if batch is None:\n",
-    "                    break\n",
-    "                # Skip incomplete last batch\n",
-    "                if batch.shape[0] == dataloader.batch_size:\n",
-    "                    superbatch.append(batch.squeeze(1))\n",
-    "        except StopIteration:\n",
-    "            pass\n",
-    "        superbatch = torch.stack(superbatch, axis=0)\n",
-    "        yield superbatch"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 170,
-   "id": "2210705b",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import os\n",
-    "\n",
-    "def encode_captioned_dataset(dataset, output_tsv, batch_size=32, num_workers=16):\n",
-    "    if os.path.isfile(output_tsv):\n",
-    "        print(f\"Destination file {output_tsv} already exists, please move away.\")\n",
-    "        return\n",
-    "    \n",
-    "    num_tpus = 8    \n",
-    "    dataloader = DataLoader(dataset, batch_size=batch_size, num_workers=num_workers)\n",
-    "    superbatches = superbatch_generator(dataloader, num_tpus=num_tpus)\n",
-    "    \n",
-    "    p_encoder = pmap(lambda batch: encode(model, batch))\n",
-    "\n",
-    "    # We save each superbatch to avoid reallocation of buffers as we process them.\n",
-    "    # We keep the file open to prevent excessive file seeks.\n",
-    "    with open(output_tsv, \"w\") as file:\n",
-    "        iterations = len(dataset) // (batch_size * num_tpus)\n",
-    "        for n in tqdm(range(iterations)):\n",
-    "            superbatch = next(superbatches)\n",
-    "            encoded = p_encoder(superbatch.numpy())\n",
-    "            encoded = encoded.reshape(-1, encoded.shape[-1])\n",
-    "\n",
-    "            # Extract fields from the dataset internal `captions` property, and save to disk\n",
-    "            start_index = n * batch_size * num_tpus\n",
-    "            end_index = (n+1) * batch_size * num_tpus\n",
-    "            paths = dataset.captions[\"image_file\"][start_index:end_index].values\n",
-    "            captions = dataset.captions[\"caption\"][start_index:end_index].values\n",
-    "            encoded_as_string = list(map(lambda item: np.array2string(item, separator=',', max_line_width=50000, formatter={'int':lambda x: str(x)}), encoded))\n",
-    "            batch_df = pd.DataFrame.from_dict({\"image_file\": paths, \"caption\": captions, \"encoding\": encoded_as_string})\n",
-    "            batch_df.to_csv(file, sep='\\t', header=(n==0), index=None)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 171,
-   "id": "7704863d",
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4850/4850 [2:27:51<00:00,  1.83s/it]\n"
-     ]
-    }
-   ],
-   "source": [
-    "encode_captioned_dataset(dataset, yfcc100m_output, batch_size=64, num_workers=16)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "8953dd84",
-   "metadata": {},
-   "source": [
-    "----"
-   ]
-  }
- ],
- "metadata": {
-  "interpreter": {
-   "hash": "db471c52d602b4f5f40ecaf278e88ccfef85c29d0a1a07185b0d51fc7acf4e26"
-  },
-  "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.8.10"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 5
-}

dev/environment.yaml

@@ -1,10 +0,0 @@
-name: dalle
-channels:
-  - defaults
-dependencies:
-  - python=3.9.5
-  - pip=21.1.3
-  - ipython=7.22.0
-  - cudatoolkit
-  - pip:
-    - -r requirements.txt

dev/requirements.txt

@@ -1,14 +0,0 @@
-requests
--f https://storage.googleapis.com/jax-releases/libtpu_releases.html
-jax[tpu]>=0.2.16
-transformers
-datasets
-flax
-jupyter
-wandb
-nltk
-optax
-git+https://github.com/patil-suraj/vqgan-jax.git@610d842dd33c739325a944102ed33acc07692dd5
-
-# Inference
-ftfy

dev/seq2seq/do_big_run.sh

@@ -1,21 +0,0 @@
-python run_seq2seq_flax.py \
-	--dataset_repo_or_path dalle-mini/encoded \
-	--train_file **/train/*/*.jsonl \
-	--validation_file **/valid/*/*.jsonl \
-	--len_train 129847128 \
-	--len_eval 157312 \
-	--eval_steps 1000 \
-	--streaming \
-	--normalize_text \
-	--output_dir output \
-	--per_device_train_batch_size 56 \
-	--per_device_eval_batch_size 56 \
-	--preprocessing_num_workers 80 \
-	--warmup_steps 5000 \
-	--gradient_accumulation_steps 8 \
-	--do_train \
-	--do_eval \
-	--adafactor \
-	--num_train_epochs 6 \
-	--log_model \
-	--learning_rate 0.005

dev/seq2seq/do_small_run.sh

@@ -1,19 +0,0 @@
-python run_seq2seq_flax.py \
-	--dataset_repo_or_path dalle-mini/encoded \
-	--train_file **/train/CC3M/*.jsonl \
-	--validation_file **/valid/*/*.jsonl \
-	--len_train 129847128 \
-	--len_eval 157312 \
-	--streaming \
-	--output_dir output \
-	--per_device_train_batch_size 16 \
-	--per_device_eval_batch_size 16 \
-	--preprocessing_num_workers 80 \
-	--warmup_steps 125 \
-	--gradient_accumulation_steps 8 \
-	--do_train \
-	--do_eval \
-	--adafactor \
-	--num_train_epochs 1 \
-	--max_train_samples 10000 \
-	--learning_rate 0.005

pyproject.toml

@@ -0,0 +1,2 @@
+[tool.isort]
+profile = "black"

setup.cfg

@@ -16,3 +16,11 @@ install_requires =
     ftfy
     jax
     flax
+
+[options.extras_require]
+dev =
+    tqdm
+    wandb
+    optax
+    black[jupyter]
+    isort

setup.py

@@ -1,4 +1,4 @@
 from setuptools import setup
 
 if __name__ == "__main__":
-    setup()
+    setup()

dev/encoding/vqgan-jax-encoding-webdataset.ipynb → tools/dataset/encode_dataset.ipynb

@@ -5,7 +5,7 @@
    "id": "d0b72877",
    "metadata": {},
    "source": [
-    "# VQGAN JAX Encoding for `webdataset`"
+    "# Pre-encoding a dataset for DALLE·mini"
    ]
   },
   {
@@ -15,7 +15,11 @@
    "source": [
     "This notebook shows how to pre-encode images to token sequences using JAX, VQGAN and a dataset in the [`webdataset` format](https://webdataset.github.io/webdataset/).\n",
     "\n",
-    "This example uses a small subset of YFCC100M we created for testing, but it should be easy to adapt to any other image/caption dataset in the `webdataset` format."
+    "Adapt it to your own dataset and image encoder.\n",
+    "\n",
+    "At the end you should have a dataset of pairs:\n",
+    "* a caption defined as a string\n",
+    "* an encoded image defined as a list of int."
    ]
   },
   {
@@ -25,19 +29,15 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import numpy as np\n",
-    "from tqdm import tqdm\n",
+    "from tqdm.notebook import tqdm\n",
     "\n",
-    "import torch\n",
     "import torchvision.transforms as T\n",
-    "import torchvision.transforms.functional as TF\n",
-    "from torchvision.transforms import InterpolationMode\n",
-    "import math\n",
     "\n",
     "import webdataset as wds\n",
     "\n",
     "import jax\n",
-    "from jax import pmap"
+    "import braceexpand\n",
+    "from pathlib import Path"
    ]
   },
   {
@@ -45,184 +45,110 @@
    "id": "c7c4c1e6",
    "metadata": {},
    "source": [
-    "## Dataset and Parameters"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "9822850f",
-   "metadata": {},
-   "source": [
-    "The following is the list of shards we'll process. We hardcode the length of data so that we can see nice progress bars using `tqdm`."
+    "## Configuration Parameters"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 3,
    "id": "1265dbfe",
    "metadata": {},
    "outputs": [],
    "source": [
-    "shards = 'https://huggingface.co/datasets/dalle-mini/YFCC100M_OpenAI_subset/resolve/main/data/shard-{0000..0008}.tar'\n",
-    "length = 8320"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "7e38fa14",
-   "metadata": {},
-   "source": [
-    "If we are extra cautious or our server is unreliable, we can enable retries by providing a custom `curl` retrieval command:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "4c8c5960",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# Enable curl retries to try to work around temporary network / server errors.\n",
-    "# This shouldn't be necessary when using reliable servers.\n",
-    "# shards = f'pipe:curl -s --retry 5 --retry-delay 5 -L {shards} || true'"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "13c6631b",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from pathlib import Path\n",
+    "shards = \"my_images/shard-{0000..0008}.tar\"  # defined using braceexpand format as used by webdataset\n",
+    "encoded_output = Path(\"encoded_data\")  # where we will save our encoded data\n",
     "\n",
-    "# Output directory for encoded files\n",
-    "encoded_output = Path.home()/'data'/'wds'/'encoded'\n",
+    "VQGAN_REPO, VQGAN_COMMIT_ID = (\n",
+    "    \"dalle-mini/vqgan_imagenet_f16_16384\",\n",
+    "    \"85eb5d3b51a1c62a0cc8f4ccdee9882c0d0bd384\",\n",
+    ")\n",
     "\n",
-    "batch_size = 128           # Per device\n",
-    "num_workers = 8            # For parallel processing"
+    "# good defaults for a TPU v3-8\n",
+    "batch_size = 128  # Per device\n",
+    "num_workers = 8  # For parallel processing\n",
+    "total_bs = batch_size * jax.device_count()  # You can use a smaller size while testing\n",
+    "save_frequency = 128  # Number of batches to create a new file (180MB for f16 and 720MB for f8 per file)"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
-   "id": "3435fb85",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "bs = batch_size * jax.device_count()    # You can use a smaller size while testing\n",
-    "batches = math.ceil(length / bs)"
+   "execution_count": 5,
+   "id": "cd956ec6-7d98-4d4d-a454-f80fe857eadd",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "['XXX/shard-0000.tar',\n",
+       " 'XXX/shard-0001.tar',\n",
+       " 'XXX/shard-0002.tar',\n",
+       " 'XXX/shard-0003.tar',\n",
+       " 'XXX/shard-0004.tar',\n",
+       " 'XXX/shard-0005.tar',\n",
+       " 'XXX/shard-0006.tar',\n",
+       " 'XXX/shard-0007.tar',\n",
+       " 'XXX/shard-0008.tar']"
+      ]
+     },
+     "execution_count": 5,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "shards = list(\n",
+    "    braceexpand.braceexpand(shards)\n",
+    ")  # better display for tqdm with known length"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "88598e4b",
+   "id": "75dba8e2",
    "metadata": {},
    "source": [
-    "Image processing"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "669b35df",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def center_crop(image, max_size=256):\n",
-    "    # Note: we allow upscaling too. We should exclude small images.    \n",
-    "    image = TF.resize(image, max_size, interpolation=InterpolationMode.LANCZOS)\n",
-    "    image = TF.center_crop(image, output_size=2 * [max_size])\n",
-    "    return image\n",
-    "\n",
-    "preprocess_image = T.Compose([\n",
-    "    center_crop,\n",
-    "    T.ToTensor(),\n",
-    "    lambda t: t.permute(1, 2, 0)   # Reorder, we need dimensions last\n",
-    "])"
+    "## Load data"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "a185e90c",
+   "id": "a1e8fb95",
    "metadata": {},
    "source": [
-    "Caption preparation.\n",
-    "\n",
-    "Note that we receive the contents of the `json` structure, which will be replaced by the string we return.\n",
-    "If we want to keep other fields inside `json`, we can add `caption` as a new field."
+    "We load data using `webdataset`."
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "423ee10e",
+   "id": "9ef5de9e",
    "metadata": {},
    "outputs": [],
    "source": [
-    "def create_caption(item):\n",
-    "    title = item['title_clean'].strip()\n",
-    "    description = item['description_clean'].strip()\n",
-    "    if len(title) > 0 and title[-1] not in '.!?': title += '.'\n",
-    "    return f'{title} {description}'"
+    "ds = (\n",
+    "    wds.WebDataset(shards, handler=wds.warn_and_continue)\n",
+    "    .decode(\"rgb\", handler=wds.warn_and_continue)\n",
+    "    .to_tuple(\"jpg\", \"txt\")  # assumes image is in `jpg` and caption in `txt`\n",
+    "    .batched(total_bs)  # load in batch per worker (faster)\n",
+    ")"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "8d3a95db",
+   "id": "90981824",
    "metadata": {},
    "source": [
-    "When an error occurs (a download is disconnected, an image cannot be decoded, etc) the process stops with an exception. We can use one of the exception handlers provided by the `webdataset` library, such as `wds.warn_and_continue` or `wds.ignore_and_continue` to ignore the offending entry and keep iterating.\n",
-    "\n",
-    "**IMPORTANT WARNING:** Do not use error handlers to ignore exceptions until you have tested that your processing pipeline works fine. Otherwise, the process will continue trying to find a valid entry, and it will consume your whole dataset without doing any work.\n",
-    "\n",
-    "We can also create our custom exception handler as demonstrated here:"
+    "Note:\n",
+    "* you can also shuffle shards and items using `shardshuffle` and `shuffle` if necessary.\n",
+    "* you may need to resize images in your pipeline (with `map_dict` for example), we assume they are already set to 256x256.\n",
+    "* you can also filter out some items using `select`."
    ]
   },
   {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "369d9719",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# UNUSED - Log exceptions to a file\n",
-    "def ignore_and_log(exn):\n",
-    "    with open('errors.txt', 'a') as f:\n",
-    "        f.write(f'{repr(exn)}\\n')\n",
-    "    return True"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "27de1414",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# Or simply use `wds.ignore_and_continue`\n",
-    "exception_handler = wds.warn_and_continue"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "5149b6d5",
+   "cell_type": "markdown",
+   "id": "129c377d",
    "metadata": {},
-   "outputs": [],
    "source": [
-    "dataset = wds.WebDataset(shards,\n",
-    "                         length=batches,              # Hint so `len` is implemented\n",
-    "                         shardshuffle=False,          # Keep same order for encoded files for easier bookkeeping. Set to `True` for training.\n",
-    "                         handler=exception_handler,   # Ignore read errors instead of failing.\n",
-    ")\n",
-    "\n",
-    "dataset = (dataset           \n",
-    "      .decode('pil')                     # decode image with PIL\n",
-    "#       .map_dict(jpg=preprocess_image, json=create_caption, handler=exception_handler)    # Process fields with functions defined above\n",
-    "      .map_dict(jpg=preprocess_image, json=create_caption)    # Process fields with functions defined above\n",
-    "      .to_tuple('__key__', 'jpg', 'json') # filter to keep only key (for reference), image, caption.\n",
-    "      .batched(bs))                      # better to batch in the dataset (but we could also do it in the dataloader) - this arg does not affect speed and we could remove it"
+    "We can now inspect our data."
    ]
   },
   {
@@ -235,7 +161,7 @@
    "outputs": [],
    "source": [
     "%%time\n",
-    "keys, images, captions = next(iter(dataset))"
+    "images, captions = next(iter(ds))"
    ]
   },
   {
@@ -251,54 +177,50 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "c24693c0",
+   "id": "5acfc4d8",
    "metadata": {},
    "outputs": [],
    "source": [
-    "T.ToPILImage()(images[0].permute(2, 0, 1))"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "44d50a51",
-   "metadata": {},
-   "source": [
-    "### Torch DataLoader"
+    "captions[:10]"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "e2df5e13",
+   "id": "c24693c0",
    "metadata": {},
    "outputs": [],
    "source": [
-    "dl = torch.utils.data.DataLoader(dataset, batch_size=None, num_workers=num_workers)"
+    "T.ToPILImage()(images[0].permute(2, 0, 1))"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "a354472b",
+   "id": "3059ffb1",
    "metadata": {},
    "source": [
-    "## VQGAN-JAX model"
+    "Finally we create our dataloader."
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "2fcf01d7",
+   "id": "c227c551",
    "metadata": {},
    "outputs": [],
    "source": [
-    "from vqgan_jax.modeling_flax_vqgan import VQModel"
+    "dl = (\n",
+    "    wds.WebLoader(ds, batch_size=None, num_workers=8).unbatched().batched(total_bs)\n",
+    ")  # avoid partial batch at the end of each worker"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "9daa636d",
+   "id": "a354472b",
    "metadata": {},
    "source": [
+    "## Image encoder\n",
+    "\n",
     "We'll use a VQGAN trained with Taming Transformers and converted to a JAX model."
    ]
   },
@@ -311,7 +233,11 @@
    },
    "outputs": [],
    "source": [
-    "model = VQModel.from_pretrained(\"flax-community/vqgan_f16_16384\")"
+    "from vqgan_jax.modeling_flax_vqgan import VQModel\n",
+    "from flax.jax_utils import replicate\n",
+    "\n",
+    "vqgan = VQModel.from_pretrained(\"flax-community/vqgan_f16_16384\")\n",
+    "vqgan_params = replicate(vqgan.params)"
    ]
   },
   {
@@ -327,18 +253,7 @@
    "id": "20357f74",
    "metadata": {},
    "source": [
-    "Encoding is really simple using `shard` to automatically distribute \"superbatches\" across devices, and `pmap`. This is all it takes to create our encoding function, that will be jitted on first use."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "6686b004",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from flax.training.common_utils import shard\n",
-    "from functools import partial"
+    "Encoding is really simple using `shard` to automatically distribute batches across devices and `pmap`."
    ]
   },
   {
@@ -348,21 +263,17 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "from flax.training.common_utils import shard\n",
+    "from functools import partial\n",
+    "\n",
+    "\n",
     "@partial(jax.pmap, axis_name=\"batch\")\n",
-    "def encode(batch):\n",
+    "def p_encode(batch, params):\n",
     "    # Not sure if we should `replicate` params, does not seem to have any effect\n",
-    "    _, indices = model.encode(batch)\n",
+    "    _, indices = vqgan.encode(batch, params=params)\n",
     "    return indices"
    ]
   },
-  {
-   "cell_type": "markdown",
-   "id": "14375a41",
-   "metadata": {},
-   "source": [
-    "### Encoding loop"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -370,49 +281,48 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import os\n",
     "import pandas as pd\n",
     "\n",
-    "def encode_captioned_dataset(dataloader, output_dir, save_every=14):\n",
-    "    output_dir.mkdir(parents=True, exist_ok=True)\n",
     "\n",
-    "    # Saving strategy:\n",
-    "    # - Create a new file every so often to prevent excessive file seeking.\n",
-    "    # - Save each batch after processing.\n",
-    "    # - Keep the file open until we are done with it.\n",
-    "    file = None        \n",
-    "    for n, (keys, images, captions) in enumerate(tqdm(dataloader)):\n",
-    "        if (n % save_every == 0):\n",
-    "            if file is not None:\n",
-    "                file.close()\n",
-    "            split_num = n // save_every\n",
-    "            file = open(output_dir/f'split_{split_num:05x}.jsonl', 'w')\n",
-    "\n",
-    "        images = shard(images.numpy().squeeze())\n",
-    "        encoded = encode(images)\n",
+    "def encode_dataset(dataloader, output_dir, save_frequency):\n",
+    "    output_dir.mkdir(parents=True, exist_ok=True)\n",
+    "    all_captions = []\n",
+    "    all_encoding = []\n",
+    "    n_file = 1\n",
+    "    for idx, (images, captions) in enumerate(tqdm(dataloader)):\n",
+    "        images = images.numpy()\n",
+    "        n = len(images) // 8 * 8\n",
+    "        if n != len(images):\n",
+    "            # get the max number of images we can (multiple of 8)\n",
+    "            print(f\"Different sizes {n} vs {len(images)}\")\n",
+    "            images = images[:n]\n",
+    "            captions = captions[:n]\n",
+    "        if not len(captions):\n",
+    "            print(f\"No images/captions in batch...\")\n",
+    "            continue\n",
+    "        images = shard(images)\n",
+    "        encoded = p_encode(images, vqgan_params)\n",
     "        encoded = encoded.reshape(-1, encoded.shape[-1])\n",
+    "        all_captions.extend(captions)\n",
+    "        all_encoding.extend(encoded.tolist())\n",
     "\n",
-    "        encoded_as_string = list(map(lambda item: np.array2string(item, separator=',', max_line_width=50000, formatter={'int':lambda x: str(x)}), encoded))\n",
-    "        batch_df = pd.DataFrame.from_dict({\"key\": keys, \"caption\": captions, \"encoding\": encoded_as_string})\n",
-    "        batch_df.to_json(file, orient='records', lines=True)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "09ff75a3",
-   "metadata": {},
-   "source": [
-    "Create a new file every 318 iterations. This should produce splits of ~500 MB each, when using a total batch size of 1024."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "96222bb4",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "save_every = 318"
+    "        # save files\n",
+    "        if (idx + 1) % save_frequency == 0:\n",
+    "            print(f\"Saving file {n_file}\")\n",
+    "            batch_df = pd.DataFrame.from_dict(\n",
+    "                {\"caption\": all_captions, \"encoding\": all_encoding}\n",
+    "            )\n",
+    "            batch_df.to_parquet(f\"{output_dir}/{n_file:03d}.parquet\")\n",
+    "            all_captions = []\n",
+    "            all_encoding = []\n",
+    "            n_file += 1\n",
+    "\n",
+    "    if len(all_captions):\n",
+    "        print(f\"Saving final file {n_file}\")\n",
+    "        batch_df = pd.DataFrame.from_dict(\n",
+    "            {\"caption\": all_captions, \"encoding\": all_encoding}\n",
+    "        )\n",
+    "        batch_df.to_parquet(f\"{output_dir}/{n_file:03d}.parquet\")"
    ]
   },
   {
@@ -422,7 +332,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "encode_captioned_dataset(dl, encoded_output, save_every=save_every)"
+    "encode_dataset(dl, output_dir=encoded_output, save_frequency=save_frequency)"
    ]
   },
   {
@@ -453,7 +363,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.10"
+   "version": "3.9.7"
   }
  },
  "nbformat": 4,

tools/inference/log_inference_samples.ipynb

@@ -31,11 +31,14 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "run_ids = ['63otg87g']\n",
-    "ENTITY, PROJECT = 'dalle-mini', 'dalle-mini'  # used only for training run\n",
-    "VQGAN_REPO, VQGAN_COMMIT_ID = 'dalle-mini/vqgan_imagenet_f16_16384', 'e93a26e7707683d349bf5d5c41c5b0ef69b677a9'\n",
-    "latest_only = True    # log only latest or all versions\n",
-    "suffix = ''           # mainly for duplicate inference runs with a deleted version\n",
+    "run_ids = [\"63otg87g\"]\n",
+    "ENTITY, PROJECT = \"dalle-mini\", \"dalle-mini\"  # used only for training run\n",
+    "VQGAN_REPO, VQGAN_COMMIT_ID = (\n",
+    "    \"dalle-mini/vqgan_imagenet_f16_16384\",\n",
+    "    \"e93a26e7707683d349bf5d5c41c5b0ef69b677a9\",\n",
+    ")\n",
+    "latest_only = True  # log only latest or all versions\n",
+    "suffix = \"\"  # mainly for duplicate inference runs with a deleted version\n",
     "add_clip_32 = False"
    ]
   },
@@ -63,8 +66,8 @@
     "num_images = 128\n",
     "top_k = 8\n",
     "text_normalizer = TextNormalizer()\n",
-    "padding_item = 'NONE'\n",
-    "seed = random.randint(0, 2**32-1)\n",
+    "padding_item = \"NONE\"\n",
+    "seed = random.randint(0, 2 ** 32 - 1)\n",
     "key = jax.random.PRNGKey(seed)\n",
     "api = wandb.Api()"
    ]
@@ -100,12 +103,15 @@
     "def p_decode(indices, params):\n",
     "    return vqgan.decode_code(indices, params=params)\n",
     "\n",
+    "\n",
     "@partial(jax.pmap, axis_name=\"batch\")\n",
     "def p_clip16(inputs, params):\n",
     "    logits = clip16(params=params, **inputs).logits_per_image\n",
     "    return logits\n",
     "\n",
+    "\n",
     "if add_clip_32:\n",
+    "\n",
     "    @partial(jax.pmap, axis_name=\"batch\")\n",
     "    def p_clip32(inputs, params):\n",
     "        logits = clip32(params=params, **inputs).logits_per_image\n",
@@ -119,13 +125,13 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "with open('samples.txt', encoding='utf8') as f:\n",
+    "with open(\"samples.txt\", encoding=\"utf8\") as f:\n",
     "    samples = [l.strip() for l in f.readlines()]\n",
     "    # make list multiple of batch_size by adding elements\n",
     "    samples_to_add = [padding_item] * (-len(samples) % batch_size)\n",
     "    samples.extend(samples_to_add)\n",
     "    # reshape\n",
-    "    samples = [samples[i:i+batch_size] for i in range(0, len(samples), batch_size)]"
+    "    samples = [samples[i : i + batch_size] for i in range(0, len(samples), batch_size)]"
    ]
   },
   {
@@ -138,9 +144,17 @@
     "def get_artifact_versions(run_id, latest_only=False):\n",
     "    try:\n",
     "        if latest_only:\n",
-    "            return [api.artifact(type='bart_model', name=f'{ENTITY}/{PROJECT}/model-{run_id}:latest')]\n",
+    "            return [\n",
+    "                api.artifact(\n",
+    "                    type=\"bart_model\", name=f\"{ENTITY}/{PROJECT}/model-{run_id}:latest\"\n",
+    "                )\n",
+    "            ]\n",
     "        else:\n",
-    "            return api.artifact_versions(type_name='bart_model', name=f'{ENTITY}/{PROJECT}/model-{run_id}', per_page=10000)\n",
+    "            return api.artifact_versions(\n",
+    "                type_name=\"bart_model\",\n",
+    "                name=f\"{ENTITY}/{PROJECT}/model-{run_id}\",\n",
+    "                per_page=10000,\n",
+    "            )\n",
     "    except:\n",
     "        return []"
    ]
@@ -153,7 +167,7 @@
    "outputs": [],
    "source": [
     "def get_training_config(run_id):\n",
-    "    training_run = api.run(f'{ENTITY}/{PROJECT}/{run_id}')\n",
+    "    training_run = api.run(f\"{ENTITY}/{PROJECT}/{run_id}\")\n",
     "    config = training_run.config\n",
     "    return config"
    ]
@@ -168,8 +182,8 @@
     "# retrieve inference run details\n",
     "def get_last_inference_version(run_id):\n",
     "    try:\n",
-    "        inference_run = api.run(f'dalle-mini/dalle-mini/{run_id}-clip16{suffix}')\n",
-    "        return inference_run.summary.get('version', None)\n",
+    "        inference_run = api.run(f\"dalle-mini/dalle-mini/{run_id}-clip16{suffix}\")\n",
+    "        return inference_run.summary.get(\"version\", None)\n",
     "    except:\n",
     "        return None"
    ]
@@ -183,7 +197,6 @@
    "source": [
     "# compile functions - needed only once per run\n",
     "def pmap_model_function(model):\n",
-    "    \n",
     "    @partial(jax.pmap, axis_name=\"batch\")\n",
     "    def _generate(tokenized_prompt, key, params):\n",
     "        return model.generate(\n",
@@ -195,7 +208,7 @@
     "            top_k=gen_top_k,\n",
     "            top_p=gen_top_p\n",
     "        )\n",
-    "    \n",
+    "\n",
     "    return _generate"
    ]
   },
@@ -222,13 +235,21 @@
     "training_config = get_training_config(run_id)\n",
     "run = None\n",
     "p_generate = None\n",
-    "model_files = ['config.json', 'flax_model.msgpack', 'merges.txt', 'special_tokens_map.json', 'tokenizer.json', 'tokenizer_config.json', 'vocab.json']\n",
+    "model_files = [\n",
+    "    \"config.json\",\n",
+    "    \"flax_model.msgpack\",\n",
+    "    \"merges.txt\",\n",
+    "    \"special_tokens_map.json\",\n",
+    "    \"tokenizer.json\",\n",
+    "    \"tokenizer_config.json\",\n",
+    "    \"vocab.json\",\n",
+    "]\n",
     "for artifact in artifact_versions:\n",
-    "    print(f'Processing artifact: {artifact.name}')\n",
+    "    print(f\"Processing artifact: {artifact.name}\")\n",
     "    version = int(artifact.version[1:])\n",
     "    results16, results32 = [], []\n",
-    "    columns = ['Caption'] + [f'Image {i+1}' for i in range(top_k)]\n",
-    "    \n",
+    "    columns = [\"Caption\"] + [f\"Image {i+1}\" for i in range(top_k)]\n",
+    "\n",
     "    if latest_only:\n",
     "        assert last_inference_version is None or version > last_inference_version\n",
     "    else:\n",
@@ -236,14 +257,23 @@
     "            # we should start from v0\n",
     "            assert version == 0\n",
     "        elif version <= last_inference_version:\n",
-    "            print(f'v{version} has already been logged (versions logged up to v{last_inference_version}')\n",
+    "            print(\n",
+    "                f\"v{version} has already been logged (versions logged up to v{last_inference_version}\"\n",
+    "            )\n",
     "        else:\n",
     "            # check we are logging the correct version\n",
     "            assert version == last_inference_version + 1\n",
     "\n",
     "    # start/resume corresponding run\n",
     "    if run is None:\n",
-    "        run = wandb.init(job_type='inference', entity='dalle-mini', project='dalle-mini', config=training_config, id=f'{run_id}-clip16{suffix}', resume='allow')\n",
+    "        run = wandb.init(\n",
+    "            job_type=\"inference\",\n",
+    "            entity=\"dalle-mini\",\n",
+    "            project=\"dalle-mini\",\n",
+    "            config=training_config,\n",
+    "            id=f\"{run_id}-clip16{suffix}\",\n",
+    "            resume=\"allow\",\n",
+    "        )\n",
     "\n",
     "    # work in temporary directory\n",
     "    with tempfile.TemporaryDirectory() as tmp:\n",
@@ -264,64 +294,109 @@
     "\n",
     "        # process one batch of captions\n",
     "        for batch in tqdm(samples):\n",
-    "            processed_prompts = [text_normalizer(x) for x in batch] if model.config.normalize_text else list(batch)\n",
+    "            processed_prompts = (\n",
+    "                [text_normalizer(x) for x in batch]\n",
+    "                if model.config.normalize_text\n",
+    "                else list(batch)\n",
+    "            )\n",
     "\n",
     "            # repeat the prompts to distribute over each device and tokenize\n",
     "            processed_prompts = processed_prompts * jax.device_count()\n",
-    "            tokenized_prompt = tokenizer(processed_prompts, return_tensors='jax', padding='max_length', truncation=True, max_length=128).data\n",
+    "            tokenized_prompt = tokenizer(\n",
+    "                processed_prompts,\n",
+    "                return_tensors=\"jax\",\n",
+    "                padding=\"max_length\",\n",
+    "                truncation=True,\n",
+    "                max_length=128,\n",
+    "            ).data\n",
     "            tokenized_prompt = shard(tokenized_prompt)\n",
     "\n",
     "            # generate images\n",
     "            images = []\n",
-    "            pbar = tqdm(range(num_images // jax.device_count()), desc='Generating Images', leave=True)\n",
+    "            pbar = tqdm(\n",
+    "                range(num_images // jax.device_count()),\n",
+    "                desc=\"Generating Images\",\n",
+    "                leave=True,\n",
+    "            )\n",
     "            for i in pbar:\n",
     "                key, subkey = jax.random.split(key)\n",
-    "                encoded_images = p_generate(tokenized_prompt, shard_prng_key(subkey), model_params)\n",
+    "                encoded_images = p_generate(\n",
+    "                    tokenized_prompt, shard_prng_key(subkey), model_params\n",
+    "                )\n",
     "                encoded_images = encoded_images.sequences[..., 1:]\n",
     "                decoded_images = p_decode(encoded_images, vqgan_params)\n",
-    "                decoded_images = decoded_images.clip(0., 1.).reshape((-1, 256, 256, 3))\n",
+    "                decoded_images = decoded_images.clip(0.0, 1.0).reshape(\n",
+    "                    (-1, 256, 256, 3)\n",
+    "                )\n",
     "                for img in decoded_images:\n",
-    "                    images.append(Image.fromarray(np.asarray(img * 255, dtype=np.uint8)))\n",
+    "                    images.append(\n",
+    "                        Image.fromarray(np.asarray(img * 255, dtype=np.uint8))\n",
+    "                    )\n",
     "\n",
-    "            def add_clip_results(results, processor, p_clip, clip_params):  \n",
-    "                clip_inputs = processor(text=batch, images=images, return_tensors='np', padding='max_length', max_length=77, truncation=True).data\n",
+    "            def add_clip_results(results, processor, p_clip, clip_params):\n",
+    "                clip_inputs = processor(\n",
+    "                    text=batch,\n",
+    "                    images=images,\n",
+    "                    return_tensors=\"np\",\n",
+    "                    padding=\"max_length\",\n",
+    "                    max_length=77,\n",
+    "                    truncation=True,\n",
+    "                ).data\n",
     "                # each shard will have one prompt, images need to be reorganized to be associated to the correct shard\n",
-    "                images_per_prompt_indices = np.asarray(range(0, len(images), batch_size))\n",
-    "                clip_inputs['pixel_values'] = jnp.concatenate(list(clip_inputs['pixel_values'][images_per_prompt_indices + i] for i in range(batch_size)))\n",
+    "                images_per_prompt_indices = np.asarray(\n",
+    "                    range(0, len(images), batch_size)\n",
+    "                )\n",
+    "                clip_inputs[\"pixel_values\"] = jnp.concatenate(\n",
+    "                    list(\n",
+    "                        clip_inputs[\"pixel_values\"][images_per_prompt_indices + i]\n",
+    "                        for i in range(batch_size)\n",
+    "                    )\n",
+    "                )\n",
     "                clip_inputs = shard(clip_inputs)\n",
     "                logits = p_clip(clip_inputs, clip_params)\n",
     "                logits = logits.reshape(-1, num_images)\n",
     "                top_scores = logits.argsort()[:, -top_k:][..., ::-1]\n",
     "                logits = jax.device_get(logits)\n",
     "                # add to results table\n",
-    "                for i, (idx, scores, sample) in enumerate(zip(top_scores, logits, batch)):\n",
-    "                    if sample == padding_item: continue\n",
+    "                for i, (idx, scores, sample) in enumerate(\n",
+    "                    zip(top_scores, logits, batch)\n",
+    "                ):\n",
+    "                    if sample == padding_item:\n",
+    "                        continue\n",
     "                    cur_images = [images[x] for x in images_per_prompt_indices + i]\n",
-    "                    top_images = [wandb.Image(cur_images[x], caption=f'Score: {scores[x]:.2f}') for x in idx]\n",
+    "                    top_images = [\n",
+    "                        wandb.Image(cur_images[x], caption=f\"Score: {scores[x]:.2f}\")\n",
+    "                        for x in idx\n",
+    "                    ]\n",
     "                    results.append([sample] + top_images)\n",
-    "                    \n",
+    "\n",
     "            # get clip scores\n",
-    "            pbar.set_description('Calculating CLIP 16 scores')\n",
+    "            pbar.set_description(\"Calculating CLIP 16 scores\")\n",
     "            add_clip_results(results16, processor16, p_clip16, clip16_params)\n",
-    "                \n",
+    "\n",
     "            # get clip 32 scores\n",
     "            if add_clip_32:\n",
-    "                pbar.set_description('Calculating CLIP 32 scores')\n",
+    "                pbar.set_description(\"Calculating CLIP 32 scores\")\n",
     "                add_clip_results(results32, processor32, p_clip32, clip32_params)\n",
     "\n",
     "            pbar.close()\n",
     "\n",
-    "                \n",
-    "\n",
     "    # log results\n",
     "    table = wandb.Table(columns=columns, data=results16)\n",
-    "    run.log({'Samples': table, 'version': version})\n",
+    "    run.log({\"Samples\": table, \"version\": version})\n",
     "    wandb.finish()\n",
-    "    \n",
-    "    if add_clip_32:        \n",
-    "        run = wandb.init(job_type='inference', entity='dalle-mini', project='dalle-mini', config=training_config, id=f'{run_id}-clip32{suffix}', resume='allow')\n",
+    "\n",
+    "    if add_clip_32:\n",
+    "        run = wandb.init(\n",
+    "            job_type=\"inference\",\n",
+    "            entity=\"dalle-mini\",\n",
+    "            project=\"dalle-mini\",\n",
+    "            config=training_config,\n",
+    "            id=f\"{run_id}-clip32{suffix}\",\n",
+    "            resume=\"allow\",\n",
+    "        )\n",
     "        table = wandb.Table(columns=columns, data=results32)\n",
-    "        run.log({'Samples': table, 'version': version})\n",
+    "        run.log({\"Samples\": table, \"version\": version})\n",
     "        wandb.finish()\n",
     "        run = None  # ensure we don't log on this run"
    ]

dev/seq2seq/run_seq2seq_flax.py → tools/train/train.py

@@ -18,37 +18,31 @@ Fine-tuning the library models for seq2seq, text to image.
 Script adapted from run_summarization_flax.py
 """
 
-import os
+import json
 import logging
+import os
 import sys
 import time
-from dataclasses import dataclass, field
+from dataclasses import asdict, dataclass, field
 from pathlib import Path
 from typing import Callable, Optional
-import json
 
 import datasets
-from datasets import Dataset
-from tqdm import tqdm
-from dataclasses import asdict
-
 import jax
 import jax.numpy as jnp
 import optax
 import transformers
+import wandb
+from datasets import Dataset
 from flax import jax_utils, traverse_util
-from flax.serialization import from_bytes, to_bytes
 from flax.jax_utils import unreplicate
+from flax.serialization import from_bytes, to_bytes
 from flax.training import train_state
 from flax.training.common_utils import get_metrics, onehot, shard_prng_key
-from transformers import (
-    AutoTokenizer,
-    HfArgumentParser,
-)
+from tqdm import tqdm
+from transformers import AutoTokenizer, HfArgumentParser
 from transformers.models.bart.modeling_flax_bart import BartConfig
 
-import wandb
-
 from dalle_mini.data import Dataset
 from dalle_mini.model import CustomFlaxBartForConditionalGeneration
 
@@ -797,7 +791,7 @@ def main():
 
     # init variables
     last_time = time.perf_counter()
-    train_metric = None
+    train_metrics = None
 
     for epoch in epochs:
         state.replace(epoch=jax_utils.replicate(epoch))
@@ -821,20 +815,20 @@ def main():
             last_time = new_time
 
             # train step
-            state, train_metric = p_train_step(
+            state, train_metrics = p_train_step(
                 state, batch, jax_utils.replicate(delta_time)
             )
             step = unreplicate(state.step)
 
             if step % training_args.logging_steps == 0 and jax.process_index() == 0:
                 # log metrics
-                wandb_log(unreplicate(train_metric), step=step, prefix="train")
+                metrics = unreplicate(train_metrics)
                 # log state parameters
                 state_dict = {
                     k.split("_")[-1]: unreplicate(getattr(state, k))
                     for k in ["epoch", "train_time", "train_samples"]
                 }
-                wandb_log(state_dict, step=step, prefix="train")
+                wandb_log({**metrics, **state_dict}, step=step, prefix="train")
 
             eval_metrics = None
             if training_args.eval_steps and step % training_args.eval_steps == 0:
@@ -844,12 +838,12 @@ def main():
                 run_save_model(state, eval_metrics)
 
         # log final train metrics
-        if train_metric is not None:
-            train_metric = unreplicate(train_metric)
-            wandb_log(train_metric, step=step, prefix="train")
+        if train_metrics is not None:
+            train_metrics = unreplicate(train_metrics)
+            wandb_log(train_metrics, step=step, prefix="train")
 
             epochs.write(
-                f"Epoch... ({epoch + 1}/{num_epochs} | Loss: {train_metric['loss']}, Learning Rate: {train_metric['learning_rate']})"
+                f"Epoch... ({epoch + 1}/{num_epochs} | Loss: {train_metrics['loss']}, Learning Rate: {train_metrics['learning_rate']})"
             )
 
         # Final evaluation