Deploy real interactive model

README.md

@@ -9,41 +9,10 @@ app_file: app.py
 pinned: false
 ---
 
-# 🤖 RoboFlamingo Demo - REAL MODEL
+# RoboFlamingo Interactive Demo 🤖
 
-Interactive demo with **trained RoboFlamingo checkpoint**!
+Upload images and get real predictions!
 
-## What is RoboFlamingo?
+⚠️ Enable T4 GPU in Settings for real model.
 
-RoboFlamingo = OpenFlamingo + **Robot-Specific Components**:
-- ✅ LSTM Policy Head for action prediction
-- ✅ Trained on CALVIN robot manipulation dataset
-- ✅ Outputs 7-DOF robot actions (position + rotation + gripper)
-- ✅ State-of-the-art performance: 4.09 average task length
-
-## Features
-- 📸 Upload robot camera views (third-person + gripper)
-- 💬 Natural language instructions
-- 🎯 Real trained model predictions
-- 📊 7-DOF trajectory visualization
-- 🤏 Gripper command timeline
-
-## Model Details
-- **Base**: OpenFlamingo (vision-language model)
-- **Policy Head**: LSTM with MLP
-- **Training**: CALVIN dataset (34 manipulation tasks)
-- **Parameters**: 386M trainable
-- **Performance**: SOTA on CALVIN benchmark
-
-## Requirements
-⚠️ **Enable T4 GPU** in Space Settings → Hardware
-
-## Citation
-```bibtex
-@article{li2023vision,
-  title={Vision-Language Foundation Models as Effective Robot Imitators},
-  author={Li, Xinghang and others},
-  journal={arXiv preprint arXiv:2311.01378},
-  year={2023}
-}
-```
+[Paper](https://arxiv.org/abs/2311.01378) | [Code](https://github.com/RoboFlamingo/RoboFlamingo)

app.py

@@ -1,4 +1,4 @@
-"""RoboFlamingo with TRAINED Checkpoint"""
+"""RoboFlamingo Interactive Demo"""
 import gradio as gr
 import torch
 import numpy as np
@@ -7,105 +7,33 @@ import matplotlib.pyplot as plt
 from io import BytesIO
 import sys
 
-sys.path.insert(0, '/home/user/app')
+sys.path.insert(0, '/home/user/app/open_flamingo/src')
 
-print("🤖 Loading RoboFlamingo with trained checkpoint")
+print("🚀 Initializing RoboFlamingo")
 
 MODEL_LOADED = False
-device = "cpu"
 
 try:
-    from robot_flamingo.models.factory import create_model_and_transforms
-    from huggingface_hub import hf_hub_download
-    
+    print("📦 Importing...")
+    from factory import create_model_and_transforms
+
     device = "cuda" if torch.cuda.is_available() else "cpu"
     print(f"Device: {device}")
-    
-    print("📦 Creating model structure...")
-    
-    # Create model
+
     model, image_processor, tokenizer = create_model_and_transforms(
         clip_vision_encoder_path="ViT-L-14",
         clip_vision_encoder_pretrained="openai",
         lang_encoder_path="mosaicml/mpt-1b-redpajama-200b",
         tokenizer_path="mosaicml/mpt-1b-redpajama-200b",
         cross_attn_every_n_layers=4,
+        decoder_type='lstm',
     )
-    
-    print("✅ Model structure created")
-    
-    # Download checkpoint from HuggingFace
-    print("📥 Downloading trained checkpoint from robovlms/RoboFlamingo...")
-    
-    try:
-        ckpt_path = hf_hub_download(
-            repo_id="robovlms/RoboFlamingo",
-            filename="checkpoint.pth",
-            repo_type="model"
-        )
-        
-        print(f"✅ Downloaded: {ckpt_path}")
-        
-        # Load checkpoint
-        print("📥 Loading checkpoint weights...")
-        checkpoint = torch.load(ckpt_path, map_location='cpu')
-        
-        # Try different checkpoint formats
-        if 'model_state_dict' in checkpoint:
-            state_dict = checkpoint['model_state_dict']
-        elif 'state_dict' in checkpoint:
-            state_dict = checkpoint['state_dict']
-        else:
-            state_dict = checkpoint
-        
-        # Load with strict=False to handle any mismatches
-        missing, unexpected = model.load_state_dict(state_dict, strict=False)
-        
-        if len(missing) > 0:
-            print(f"⚠️  Missing keys: {len(missing)}")
-        if len(unexpected) > 0:
-            print(f"⚠️  Unexpected keys: {len(unexpected)}")
-        
-        print("✅ LOADED TRAINED CHECKPOINT!")
-        
-    except Exception as e:
-        print(f"⚠️  Checkpoint download/load failed: {e}")
-        print("   Trying alternative checkpoint...")
-        
-        try:
-            # Try the other HF repo
-            ckpt_path = hf_hub_download(
-                repo_id="hywslxh/RoboFlamingo-MPT",
-                filename="checkpoint.pth",
-                repo_type="model"
-            )
-            checkpoint = torch.load(ckpt_path, map_location='cpu')
-            
-            if 'model_state_dict' in checkpoint:
-                state_dict = checkpoint['model_state_dict']
-            elif 'state_dict' in checkpoint:
-                state_dict = checkpoint['state_dict']
-            else:
-                state_dict = checkpoint
-            
-            model.load_state_dict(state_dict, strict=False)
-            print("✅ Loaded from hywslxh/RoboFlamingo-MPT")
-            
-        except Exception as e2:
-            print(f"⚠️  Alternative also failed: {e2}")
-            print("   Using model without checkpoint")
-    
+
     model.to(device).eval()
     MODEL_LOADED = True
-    
-    print("=" * 70)
-    print("✅ MODEL READY!")
-    print("=" * 70)
-    
+    print("✅ Model loaded!")
 except Exception as e:
-    print(f"⚠️ Failed: {e}")
-    import traceback
-    traceback.print_exc()
+    print(f"⚠️ Model failed: {e}")
 
 def plot_traj(acts):
     fig = plt.figure(figsize=(10,8))
@@ -117,7 +45,7 @@ def plot_traj(acts):
     ax.scatter(x[0], y[0], z[0], c='green', s=100, label='Start')
     ax.scatter(x[-1], y[-1], z[-1], c='red', s=100, label='End')
     ax.set_xlabel('X'); ax.set_ylabel('Y'); ax.set_zlabel('Z')
-    ax.set_title('Trajectory'); ax.legend(); ax.grid()
+    ax.legend(); ax.grid()
     buf = BytesIO()
     plt.savefig(buf, format='png', dpi=100, bbox_inches='tight')
     buf.seek(0); plt.close()
@@ -141,7 +69,7 @@ def simulate(inst):
     for t in range(12):
         p = t/12
         acts.append({'timestep': t, 'delta_x': (0.05+np.random.randn()*0.01)*p,
-                    'delta_y': (0.02+np.random.randn()*0.01)*p, 
+                    'delta_y': (0.02+np.random.randn()*0.01)*p,
                     'delta_z': (-0.03+np.random.randn()*0.01)*(1-p),
                     'qw': 0.99, 'qx': 0.01, 'qy': 0.01, 'qz': 0.01})
     return acts, [0]*6+[1]*6
@@ -150,78 +78,81 @@ def predict(inst, img1, img2):
     if not inst or not inst.strip():
         return None, None, "", "❌ Enter instruction"
     if img1 is None or img2 is None:
-        return None, None, "", "❌ Upload images"
-    
+        return None, None, "", "❌ Upload both images"
+
     try:
         if isinstance(img1, np.ndarray):
             img1 = Image.fromarray(img1)
         if isinstance(img2, np.ndarray):
             img2 = Image.fromarray(img2)
-        
+
         if not MODEL_LOADED:
             acts, grip = simulate(inst)
-            status = f"⚠️ SIM\n{inst}"
+            status = f"⚠️ SIMULATION\n{inst}\nEnable GPU for real model"
         else:
-            print(f"🤖 {inst}")
             with torch.no_grad():
                 t1 = image_processor(img1).unsqueeze(0).to(device)
                 t2 = image_processor(img2).unsqueeze(0).to(device)
-                vis = torch.stack([t1, t2], dim=1).unsqueeze(2)
-                tok = tokenizer(inst, return_tensors="pt", padding=True, 
-                              truncation=True, max_length=512).to(device)
-                
-                out = model(vision_x=vis, lang_x=tok['input_ids'])
-                
+                vis = torch.stack([t1, t2], dim=1)
+                tok = tokenizer(inst, return_tensors="pt", padding=True, truncation=True).to(device)
+                out = model(vision_x=vis, lang_x=tok['input_ids'], attention_mask=tok.get('attention_mask'))
+
                 if isinstance(out, dict):
-                    a = out.get('actions')
+                    a = out.get('actions', out.get('action'))
                     g = out.get('gripper')
                 elif isinstance(out, tuple):
                     a = out[0]
                     g = out[1] if len(out)>1 else None
                 else:
-                    a = out; g = None
-                
+                    a = out
+                    g = None
+
                 if a is not None:
                     anp = a[0].cpu().numpy()
                     acts = []
                     for t, ac in enumerate(anp):
                         if len(ac)<7: ac = np.pad(ac, (0,7-len(ac)))
-                        acts.append({'timestep': t, 'delta_x': float(ac[0]), 
-                                   'delta_y': float(ac[1]), 'delta_z': float(ac[2]),
-                                   'qw': float(ac[3]), 'qx': float(ac[4]),
+                        acts.append({'timestep': t, 'delta_x': float(ac[0]), 'delta_y': float(ac[1]),
+                                   'delta_z': float(ac[2]), 'qw': float(ac[3]), 'qx': float(ac[4]),
                                    'qy': float(ac[5]), 'qz': float(ac[6])})
-                    grip = [int(x>0.5) if np.isscalar(x) else int(x[0]>0.5) 
-                           for x in (g[0].cpu().numpy() if g is not None else [0]*len(acts))]
-                    status = f"✅ MODEL\n{inst}\n{device}"
+                    grip = [int(x>0.5) if np.isscalar(x) else int(x[0]>0.5) for x in (g[0].cpu().numpy() if g is not None else [0]*len(acts))]
+                    status = f"✅ REAL MODEL\n{inst}\n{device}"
                 else:
                     acts, grip = simulate(inst)
-                    status = "⚠️ Unexpected"
-        
-        return plot_traj(acts), plot_grip(grip),                "| T | Δx | Δy | Δz |\n|--|--|--|--|\n" +                "\n".join([f"| {a['timestep']:2d} | {a['delta_x']:6.3f} | {a['delta_y']:6.3f} | {a['delta_z']:6.3f} |" 
-                          for a in acts[:12]]), status
+                    status = f"⚠️ Unexpected output\n{inst}"
+
+        traj = plot_traj(acts)
+        gp = plot_grip(grip)
+        table = "| T | Δx | Δy | Δz |\n|--|--|--|--|\n"
+        for a in acts:
+            table += f"| {a['timestep']} | {a['delta_x']:.3f} | {a['delta_y']:.3f} | {a['delta_z']:.3f} |\n"
+
+        return traj, gp, table, status
     except Exception as e:
-        print(f"❌ {e}")
-        import traceback; traceback.print_exc()
+        print(f"Error: {e}")
         acts, grip = simulate(inst)
         return plot_traj(acts), plot_grip(grip), "", f"❌ {str(e)}"
 
-mode = "🟢 LOADED" if MODEL_LOADED else "🟡 SIM"
+mode = "🟢 REAL" if MODEL_LOADED else "🟡 SIM"
 
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown(f"# 🤖 RoboFlamingo - {mode}")
+    gr.Markdown(f"# 🤖 RoboFlamingo - {mode}\n{'Real model loaded!' if MODEL_LOADED else 'Enable GPU for real model'}")
+
     with gr.Row():
         with gr.Column():
-            inst = gr.Textbox(label="Instruction", placeholder="Pick up red block", lines=3)
+            inst = gr.Textbox(label="Instruction", placeholder="Pick up the red block", lines=3)
             with gr.Row():
                 img1 = gr.Image(label="Third-Person", type="pil", height=250)
                 img2 = gr.Image(label="Gripper", type="pil", height=250)
-            btn = gr.Button("🤖 Predict", variant="primary")
+            btn = gr.Button("🚀 Predict", variant="primary", size="lg")
             st = gr.Textbox(label="Status", lines=4, interactive=False)
         with gr.Column():
             traj = gr.Image(label="Trajectory", type="pil")
             grip = gr.Image(label="Gripper", type="pil")
+
     tab = gr.Markdown()
     btn.click(predict, [inst, img1, img2], [traj, grip, tab, st])
-    gr.Markdown("[Paper](https://arxiv.org/abs/2311.01378) | [HF](https://huggingface.co/robovlms/RoboFlamingo)")
+
+    gr.Markdown(f"**Status:** {mode} | [Paper](https://arxiv.org/abs/2311.01378)")
 
 demo.launch()

open_flamingo/eval/README.md

@@ -0,0 +1,68 @@
+# OpenFlamingo Evaluation Suite
+
+This is the evaluation module of OpenFlamingo. It contains a set of utilities for evaluating multimodal models on various benchmarking datasets.
+
+*This module is a work in progress! We will be updating this README as it develops. In the meantime, if you notice an issue, please file a Bug Report or Feature Request [here](https://github.com/mlfoundations/open_flamingo/issues/new/choose).*
+
+## Supported datasets
+
+|Dataset|Task|Metric|Evaluation method|
+|-------|----|------|-----------------|
+|[COCO](https://arxiv.org/abs/1405.0312)|Captioning|CIDEr|Generation|
+|[Flickr-30K](https://aclanthology.org/Q14-1006/)|Captioning|CIDEr|Generation|
+|[VQAv2](https://arxiv.org/abs/1612.00837v3)|VQA|VQA accuracy|Generation|
+|[OK-VQA](https://arxiv.org/abs/1906.00067)|VQA|VQA accuracy|Generation|
+|[TextVQA](https://arxiv.org/abs/1904.08920)|VQA|VQA accuracy|Generation|
+|[VizWiz](https://arxiv.org/abs/1802.08218)|VQA|VQA accuracy|Generation|
+|[Hateful Memes](https://arxiv.org/abs/2005.04790)|Classification|ROC AUC|Logprobs|
+|[ImageNet](https://arxiv.org/abs/1409.0575)|Classification|Top-1 accuracy|Logprobs|
+
+When evaluating a model using `num_shots` shots, we sample the exemplars from the training split. Performance is evaluated on a disjoint test split, subsampled to `--num_samples` examples (or using the full test split if `--num_samples=-1`).
+
+## Sample scripts
+Our codebase uses DistributedDataParallel to parallelize evaluation by default, so please make sure to set the `MASTER_ADDR` and `MASTER_PORT` environment variables or use `torchrun`. We provide a sample Slurm evaluation script in `open_flamingo/open_flamingo/scripts/run_eval.sh`. 
+
+We also support evaluating at a lower precision using the `--precision` flag. We find minimal difference between evaluating at full precision vs. amp_bf16.
+
+To evaluate one of our pretrained checkpoints, we suggest first downloading a local copy of the weights, as follows:
+
+```
+# grab model checkpoint from huggingface hub
+from huggingface_hub import hf_hub_download
+HF_TOKEN="<your-hf-token-here>"
+
+checkpoint_path = hf_hub_download("openflamingo/OpenFlamingo-3B-vitl-mpt1b", "checkpoint.pt")
+checkpoint_path= hf_hub_download("openflamingo/OpenFlamingo-3B-vitl-mpt1b", 
+  "checkpoint.pt", 
+  local_dir="openflamingo/OpenFlamingo-3B-vitl-mpt1b", 
+  cache_dir="openflamingo/OpenFlamingo-3B-vitl-mpt1b", 
+  local_dir_use_symlinks=False,
+  token=HF_TOKEN)
+print(checkpoint_path)
+## openflamingo/OpenFlamingo-3B-vitl-mpt1b/checkpoint.pt
+```
+
+This should place the OpenFlamingo model at the expected location in the evaluation script.
+
+For TextVQA and VizWiz we expect annotations to be formatted differently than the original datasets. We provide the custom annotations in `open_flamingo/open_flamingo/eval/data/`. We have also uploaded all the annotation files in a [huggingface dataset](https://huggingface.co/datasets/openflamingo/eval_benchmark/tree/main) for easy access.
+
+# Evaluating using RICES (Retrieval-based In-Context Example Selection)
+
+We provide the option to evaluate using RICES, which is a method for selecting exemplars from the training set based on image similarity. This method was used in DeepMind's implementation for evaluating on ImageNet, but can be used for any dataset in our evaluation suite.
+
+To use RICES, you must first create features for a benchmark's training set. We provide a script for doing so in `open_flamingo/open_flamingo/scripts/cache_rices_features.py`. This script will extract image features for a given dataset using a given CLIP model checkpoint. For example, to extract features for the COCO training set, you can run:
+
+```bash
+python cache_rices_features.py \
+  --vision_encoder_path ViT-L-14 \
+  --vision_encoder_pretrained openai \
+  --batch_size 128 \
+  --eval_coco \
+  --coco_train_image_dir_path /path/to/coco/train2014 \
+  --coco_val_image_dir_path /path/to/coco/val2014 \
+  --coco_karpathy_json_path /path/to/coco/dataset_coco.json \
+  --coco_annotations_json_path /path/to/coco/annotations/captions_train2014.json \
+  --output_dir /path/to/coco/features 
+```
+
+This will create a directory at `/path/to/coco/features` containing a file named `coco.pkl` with the extracted features. You can then use this directory to evaluate using RICES by passing the `--rices` flag to the evaluation script, specifying the path to the features directory using the `--cached_demonstration_features` flag, and specifying the vision encoder to use for RICES using the `--rices_vision_encoder_path` and `--rices_vision_encoder_pretrained` flags.

open_flamingo/eval/__init__.py

@@ -0,0 +1 @@
+

open_flamingo/eval/classification_utils.py

@@ -0,0 +1,1008 @@
+# classnames via https://github.com/mlfoundations/wise-ft/blob/master/src/datasets/imagenet_classnames.py#L1
+IMAGENET_CLASSNAMES = [
+    "tench",
+    "goldfish",
+    "great white shark",
+    "tiger shark",
+    "hammerhead shark",
+    "electric ray",
+    "stingray",
+    "rooster",
+    "hen",
+    "ostrich",
+    "brambling",
+    "goldfinch",
+    "house finch",
+    "junco",
+    "indigo bunting",
+    "American robin",
+    "bulbul",
+    "jay",
+    "magpie",
+    "chickadee",
+    "American dipper",
+    "kite (bird of prey)",
+    "bald eagle",
+    "vulture",
+    "great grey owl",
+    "fire salamander",
+    "smooth newt",
+    "newt",
+    "spotted salamander",
+    "axolotl",
+    "American bullfrog",
+    "tree frog",
+    "tailed frog",
+    "loggerhead sea turtle",
+    "leatherback sea turtle",
+    "mud turtle",
+    "terrapin",
+    "box turtle",
+    "banded gecko",
+    "green iguana",
+    "Carolina anole",
+    "desert grassland whiptail lizard",
+    "agama",
+    "frilled-necked lizard",
+    "alligator lizard",
+    "Gila monster",
+    "European green lizard",
+    "chameleon",
+    "Komodo dragon",
+    "Nile crocodile",
+    "American alligator",
+    "triceratops",
+    "worm snake",
+    "ring-necked snake",
+    "eastern hog-nosed snake",
+    "smooth green snake",
+    "kingsnake",
+    "garter snake",
+    "water snake",
+    "vine snake",
+    "night snake",
+    "boa constrictor",
+    "African rock python",
+    "Indian cobra",
+    "green mamba",
+    "sea snake",
+    "Saharan horned viper",
+    "eastern diamondback rattlesnake",
+    "sidewinder rattlesnake",
+    "trilobite",
+    "harvestman",
+    "scorpion",
+    "yellow garden spider",
+    "barn spider",
+    "European garden spider",
+    "southern black widow",
+    "tarantula",
+    "wolf spider",
+    "tick",
+    "centipede",
+    "black grouse",
+    "ptarmigan",
+    "ruffed grouse",
+    "prairie grouse",
+    "peafowl",
+    "quail",
+    "partridge",
+    "african grey parrot",
+    "macaw",
+    "sulphur-crested cockatoo",
+    "lorikeet",
+    "coucal",
+    "bee eater",
+    "hornbill",
+    "hummingbird",
+    "jacamar",
+    "toucan",
+    "duck",
+    "red-breasted merganser",
+    "goose",
+    "black swan",
+    "tusker",
+    "echidna",
+    "platypus",
+    "wallaby",
+    "koala",
+    "wombat",
+    "jellyfish",
+    "sea anemone",
+    "brain coral",
+    "flatworm",
+    "nematode",
+    "conch",
+    "snail",
+    "slug",
+    "sea slug",
+    "chiton",
+    "chambered nautilus",
+    "Dungeness crab",
+    "rock crab",
+    "fiddler crab",
+    "red king crab",
+    "American lobster",
+    "spiny lobster",
+    "crayfish",
+    "hermit crab",
+    "isopod",
+    "white stork",
+    "black stork",
+    "spoonbill",
+    "flamingo",
+    "little blue heron",
+    "great egret",
+    "bittern bird",
+    "crane bird",
+    "limpkin",
+    "common gallinule",
+    "American coot",
+    "bustard",
+    "ruddy turnstone",
+    "dunlin",
+    "common redshank",
+    "dowitcher",
+    "oystercatcher",
+    "pelican",
+    "king penguin",
+    "albatross",
+    "grey whale",
+    "killer whale",
+    "dugong",
+    "sea lion",
+    "Chihuahua",
+    "Japanese Chin",
+    "Maltese",
+    "Pekingese",
+    "Shih Tzu",
+    "King Charles Spaniel",
+    "Papillon",
+    "toy terrier",
+    "Rhodesian Ridgeback",
+    "Afghan Hound",
+    "Basset Hound",
+    "Beagle",
+    "Bloodhound",
+    "Bluetick Coonhound",
+    "Black and Tan Coonhound",
+    "Treeing Walker Coonhound",
+    "English foxhound",
+    "Redbone Coonhound",
+    "borzoi",
+    "Irish Wolfhound",
+    "Italian Greyhound",
+    "Whippet",
+    "Ibizan Hound",
+    "Norwegian Elkhound",
+    "Otterhound",
+    "Saluki",
+    "Scottish Deerhound",
+    "Weimaraner",
+    "Staffordshire Bull Terrier",
+    "American Staffordshire Terrier",
+    "Bedlington Terrier",
+    "Border Terrier",
+    "Kerry Blue Terrier",
+    "Irish Terrier",
+    "Norfolk Terrier",
+    "Norwich Terrier",
+    "Yorkshire Terrier",
+    "Wire Fox Terrier",
+    "Lakeland Terrier",
+    "Sealyham Terrier",
+    "Airedale Terrier",
+    "Cairn Terrier",
+    "Australian Terrier",
+    "Dandie Dinmont Terrier",
+    "Boston Terrier",
+    "Miniature Schnauzer",
+    "Giant Schnauzer",
+    "Standard Schnauzer",
+    "Scottish Terrier",
+    "Tibetan Terrier",
+    "Australian Silky Terrier",
+    "Soft-coated Wheaten Terrier",
+    "West Highland White Terrier",
+    "Lhasa Apso",
+    "Flat-Coated Retriever",
+    "Curly-coated Retriever",
+    "Golden Retriever",
+    "Labrador Retriever",
+    "Chesapeake Bay Retriever",
+    "German Shorthaired Pointer",
+    "Vizsla",
+    "English Setter",
+    "Irish Setter",
+    "Gordon Setter",
+    "Brittany dog",
+    "Clumber Spaniel",
+    "English Springer Spaniel",
+    "Welsh Springer Spaniel",
+    "Cocker Spaniel",
+    "Sussex Spaniel",
+    "Irish Water Spaniel",
+    "Kuvasz",
+    "Schipperke",
+    "Groenendael dog",
+    "Malinois",
+    "Briard",
+    "Australian Kelpie",
+    "Komondor",
+    "Old English Sheepdog",
+    "Shetland Sheepdog",
+    "collie",
+    "Border Collie",
+    "Bouvier des Flandres dog",
+    "Rottweiler",
+    "German Shepherd Dog",
+    "Dobermann",
+    "Miniature Pinscher",
+    "Greater Swiss Mountain Dog",
+    "Bernese Mountain Dog",
+    "Appenzeller Sennenhund",
+    "Entlebucher Sennenhund",
+    "Boxer",
+    "Bullmastiff",
+    "Tibetan Mastiff",
+    "French Bulldog",
+    "Great Dane",
+    "St. Bernard",
+    "husky",
+    "Alaskan Malamute",
+    "Siberian Husky",
+    "Dalmatian",
+    "Affenpinscher",
+    "Basenji",
+    "pug",
+    "Leonberger",
+    "Newfoundland dog",
+    "Great Pyrenees dog",
+    "Samoyed",
+    "Pomeranian",
+    "Chow Chow",
+    "Keeshond",
+    "brussels griffon",
+    "Pembroke Welsh Corgi",
+    "Cardigan Welsh Corgi",
+    "Toy Poodle",
+    "Miniature Poodle",
+    "Standard Poodle",
+    "Mexican hairless dog (xoloitzcuintli)",
+    "grey wolf",
+    "Alaskan tundra wolf",
+    "red wolf or maned wolf",
+    "coyote",
+    "dingo",
+    "dhole",
+    "African wild dog",
+    "hyena",
+    "red fox",
+    "kit fox",
+    "Arctic fox",
+    "grey fox",
+    "tabby cat",
+    "tiger cat",
+    "Persian cat",
+    "Siamese cat",
+    "Egyptian Mau",
+    "cougar",
+    "lynx",
+    "leopard",
+    "snow leopard",
+    "jaguar",
+    "lion",
+    "tiger",
+    "cheetah",
+    "brown bear",
+    "American black bear",
+    "polar bear",
+    "sloth bear",
+    "mongoose",
+    "meerkat",
+    "tiger beetle",
+    "ladybug",
+    "ground beetle",
+    "longhorn beetle",
+    "leaf beetle",
+    "dung beetle",
+    "rhinoceros beetle",
+    "weevil",
+    "fly",
+    "bee",
+    "ant",
+    "grasshopper",
+    "cricket insect",
+    "stick insect",
+    "cockroach",
+    "praying mantis",
+    "cicada",
+    "leafhopper",
+    "lacewing",
+    "dragonfly",
+    "damselfly",
+    "red admiral butterfly",
+    "ringlet butterfly",
+    "monarch butterfly",
+    "small white butterfly",
+    "sulphur butterfly",
+    "gossamer-winged butterfly",
+    "starfish",
+    "sea urchin",
+    "sea cucumber",
+    "cottontail rabbit",
+    "hare",
+    "Angora rabbit",
+    "hamster",
+    "porcupine",
+    "fox squirrel",
+    "marmot",
+    "beaver",
+    "guinea pig",
+    "common sorrel horse",
+    "zebra",
+    "pig",
+    "wild boar",
+    "warthog",
+    "hippopotamus",
+    "ox",
+    "water buffalo",
+    "bison",
+    "ram (adult male sheep)",
+    "bighorn sheep",
+    "Alpine ibex",
+    "hartebeest",
+    "impala (antelope)",
+    "gazelle",
+    "arabian camel",
+    "llama",
+    "weasel",
+    "mink",
+    "European polecat",
+    "black-footed ferret",
+    "otter",
+    "skunk",
+    "badger",
+    "armadillo",
+    "three-toed sloth",
+    "orangutan",
+    "gorilla",
+    "chimpanzee",
+    "gibbon",
+    "siamang",
+    "guenon",
+    "patas monkey",
+    "baboon",
+    "macaque",
+    "langur",
+    "black-and-white colobus",
+    "proboscis monkey",
+    "marmoset",
+    "white-headed capuchin",
+    "howler monkey",
+    "titi monkey",
+    "Geoffroy's spider monkey",
+    "common squirrel monkey",
+    "ring-tailed lemur",
+    "indri",
+    "Asian elephant",
+    "African bush elephant",
+    "red panda",
+    "giant panda",
+    "snoek fish",
+    "eel",
+    "silver salmon",
+    "rock beauty fish",
+    "clownfish",
+    "sturgeon",
+    "gar fish",
+    "lionfish",
+    "pufferfish",
+    "abacus",
+    "abaya",
+    "academic gown",
+    "accordion",
+    "acoustic guitar",
+    "aircraft carrier",
+    "airliner",
+    "airship",
+    "altar",
+    "ambulance",
+    "amphibious vehicle",
+    "analog clock",
+    "apiary",
+    "apron",
+    "trash can",
+    "assault rifle",
+    "backpack",
+    "bakery",
+    "balance beam",
+    "balloon",
+    "ballpoint pen",
+    "Band-Aid",
+    "banjo",
+    "baluster / handrail",
+    "barbell",
+    "barber chair",
+    "barbershop",
+    "barn",
+    "barometer",
+    "barrel",
+    "wheelbarrow",
+    "baseball",
+    "basketball",
+    "bassinet",
+    "bassoon",
+    "swimming cap",
+    "bath towel",
+    "bathtub",
+    "station wagon",
+    "lighthouse",
+    "beaker",
+    "military hat (bearskin or shako)",
+    "beer bottle",
+    "beer glass",
+    "bell tower",
+    "baby bib",
+    "tandem bicycle",
+    "bikini",
+    "ring binder",
+    "binoculars",
+    "birdhouse",
+    "boathouse",
+    "bobsleigh",
+    "bolo tie",
+    "poke bonnet",
+    "bookcase",
+    "bookstore",
+    "bottle cap",
+    "hunting bow",
+    "bow tie",
+    "brass memorial plaque",
+    "bra",
+    "breakwater",
+    "breastplate",
+    "broom",
+    "bucket",
+    "buckle",
+    "bulletproof vest",
+    "high-speed train",
+    "butcher shop",
+    "taxicab",
+    "cauldron",
+    "candle",
+    "cannon",
+    "canoe",
+    "can opener",
+    "cardigan",
+    "car mirror",
+    "carousel",
+    "tool kit",
+    "cardboard box / carton",
+    "car wheel",
+    "automated teller machine",
+    "cassette",
+    "cassette player",
+    "castle",
+    "catamaran",
+    "CD player",
+    "cello",
+    "mobile phone",
+    "chain",
+    "chain-link fence",
+    "chain mail",
+    "chainsaw",
+    "storage chest",
+    "chiffonier",
+    "bell or wind chime",
+    "china cabinet",
+    "Christmas stocking",
+    "church",
+    "movie theater",
+    "cleaver",
+    "cliff dwelling",
+    "cloak",
+    "clogs",
+    "cocktail shaker",
+    "coffee mug",
+    "coffeemaker",
+    "spiral or coil",
+    "combination lock",
+    "computer keyboard",
+    "candy store",
+    "container ship",
+    "convertible",
+    "corkscrew",
+    "cornet",
+    "cowboy boot",
+    "cowboy hat",
+    "cradle",
+    "construction crane",
+    "crash helmet",
+    "crate",
+    "infant bed",
+    "Crock Pot",
+    "croquet ball",
+    "crutch",
+    "cuirass",
+    "dam",
+    "desk",
+    "desktop computer",
+    "rotary dial telephone",
+    "diaper",
+    "digital clock",
+    "digital watch",
+    "dining table",
+    "dishcloth",
+    "dishwasher",
+    "disc brake",
+    "dock",
+    "dog sled",
+    "dome",
+    "doormat",
+    "drilling rig",
+    "drum",
+    "drumstick",
+    "dumbbell",
+    "Dutch oven",
+    "electric fan",
+    "electric guitar",
+    "electric locomotive",
+    "entertainment center",
+    "envelope",
+    "espresso machine",
+    "face powder",
+    "feather boa",
+    "filing cabinet",
+    "fireboat",
+    "fire truck",
+    "fire screen",
+    "flagpole",
+    "flute",
+    "folding chair",
+    "football helmet",
+    "forklift",
+    "fountain",
+    "fountain pen",
+    "four-poster bed",
+    "freight car",
+    "French horn",
+    "frying pan",
+    "fur coat",
+    "garbage truck",
+    "gas mask or respirator",
+    "gas pump",
+    "goblet",
+    "go-kart",
+    "golf ball",
+    "golf cart",
+    "gondola",
+    "gong",
+    "gown",
+    "grand piano",
+    "greenhouse",
+    "radiator grille",
+    "grocery store",
+    "guillotine",
+    "hair clip",
+    "hair spray",
+    "half-track",
+    "hammer",
+    "hamper",
+    "hair dryer",
+    "hand-held computer",
+    "handkerchief",
+    "hard disk drive",
+    "harmonica",
+    "harp",
+    "combine harvester",
+    "hatchet",
+    "holster",
+    "home theater",
+    "honeycomb",
+    "hook",
+    "hoop skirt",
+    "gymnastic horizontal bar",
+    "horse-drawn vehicle",
+    "hourglass",
+    "iPod",
+    "clothes iron",
+    "carved pumpkin",
+    "jeans",
+    "jeep",
+    "T-shirt",
+    "jigsaw puzzle",
+    "rickshaw",
+    "joystick",
+    "kimono",
+    "knee pad",
+    "knot",
+    "lab coat",
+    "ladle",
+    "lampshade",
+    "laptop computer",
+    "lawn mower",
+    "lens cap",
+    "letter opener",
+    "library",
+    "lifeboat",
+    "lighter",
+    "limousine",
+    "ocean liner",
+    "lipstick",
+    "slip-on shoe",
+    "lotion",
+    "music speaker",
+    "loupe magnifying glass",
+    "sawmill",
+    "magnetic compass",
+    "messenger bag",
+    "mailbox",
+    "tights",
+    "one-piece bathing suit",
+    "manhole cover",
+    "maraca",
+    "marimba",
+    "mask",
+    "matchstick",
+    "maypole",
+    "maze",
+    "measuring cup",
+    "medicine cabinet",
+    "megalith",
+    "microphone",
+    "microwave oven",
+    "military uniform",
+    "milk can",
+    "minibus",
+    "miniskirt",
+    "minivan",
+    "missile",
+    "mitten",
+    "mixing bowl",
+    "mobile home",
+    "ford model t",
+    "modem",
+    "monastery",
+    "monitor",
+    "moped",
+    "mortar and pestle",
+    "graduation cap",
+    "mosque",
+    "mosquito net",
+    "vespa",
+    "mountain bike",
+    "tent",
+    "computer mouse",
+    "mousetrap",
+    "moving van",
+    "muzzle",
+    "metal nail",
+    "neck brace",
+    "necklace",
+    "baby pacifier",
+    "notebook computer",
+    "obelisk",
+    "oboe",
+    "ocarina",
+    "odometer",
+    "oil filter",
+    "pipe organ",
+    "oscilloscope",
+    "overskirt",
+    "bullock cart",
+    "oxygen mask",
+    "product packet / packaging",
+    "paddle",
+    "paddle wheel",
+    "padlock",
+    "paintbrush",
+    "pajamas",
+    "palace",
+    "pan flute",
+    "paper towel",
+    "parachute",
+    "parallel bars",
+    "park bench",
+    "parking meter",
+    "railroad car",
+    "patio",
+    "payphone",
+    "pedestal",
+    "pencil case",
+    "pencil sharpener",
+    "perfume",
+    "Petri dish",
+    "photocopier",
+    "plectrum",
+    "Pickelhaube",
+    "picket fence",
+    "pickup truck",
+    "pier",
+    "piggy bank",
+    "pill bottle",
+    "pillow",
+    "ping-pong ball",
+    "pinwheel",
+    "pirate ship",
+    "drink pitcher",
+    "block plane",
+    "planetarium",
+    "plastic bag",
+    "plate rack",
+    "farm plow",
+    "plunger",
+    "Polaroid camera",
+    "pole",
+    "police van",
+    "poncho",
+    "pool table",
+    "soda bottle",
+    "plant pot",
+    "potter's wheel",
+    "power drill",
+    "prayer rug",
+    "printer",
+    "prison",
+    "missile",
+    "projector",
+    "hockey puck",
+    "punching bag",
+    "purse",
+    "quill",
+    "quilt",
+    "race car",
+    "racket",
+    "radiator",
+    "radio",
+    "radio telescope",
+    "rain barrel",
+    "recreational vehicle",
+    "fishing casting reel",
+    "reflex camera",
+    "refrigerator",
+    "remote control",
+    "restaurant",
+    "revolver",
+    "rifle",
+    "rocking chair",
+    "rotisserie",
+    "eraser",
+    "rugby ball",
+    "ruler measuring stick",
+    "sneaker",
+    "safe",
+    "safety pin",
+    "salt shaker",
+    "sandal",
+    "sarong",
+    "saxophone",
+    "scabbard",
+    "weighing scale",
+    "school bus",
+    "schooner",
+    "scoreboard",
+    "CRT monitor",
+    "screw",
+    "screwdriver",
+    "seat belt",
+    "sewing machine",
+    "shield",
+    "shoe store",
+    "shoji screen / room divider",
+    "shopping basket",
+    "shopping cart",
+    "shovel",
+    "shower cap",
+    "shower curtain",
+    "ski",
+    "balaclava ski mask",
+    "sleeping bag",
+    "slide rule",
+    "sliding door",
+    "slot machine",
+    "snorkel",
+    "snowmobile",
+    "snowplow",
+    "soap dispenser",
+    "soccer ball",
+    "sock",
+    "solar thermal collector",
+    "sombrero",
+    "soup bowl",
+    "keyboard space bar",
+    "space heater",
+    "space shuttle",
+    "spatula",
+    "motorboat",
+    "spider web",
+    "spindle",
+    "sports car",
+    "spotlight",
+    "stage",
+    "steam locomotive",
+    "through arch bridge",
+    "steel drum",
+    "stethoscope",
+    "scarf",
+    "stone wall",
+    "stopwatch",
+    "stove",
+    "strainer",
+    "tram",
+    "stretcher",
+    "couch",
+    "stupa",
+    "submarine",
+    "suit",
+    "sundial",
+    "sunglasses",
+    "sunglasses",
+    "sunscreen",
+    "suspension bridge",
+    "mop",
+    "sweatshirt",
+    "swim trunks / shorts",
+    "swing",
+    "electrical switch",
+    "syringe",
+    "table lamp",
+    "tank",
+    "tape player",
+    "teapot",
+    "teddy bear",
+    "television",
+    "tennis ball",
+    "thatched roof",
+    "front curtain",
+    "thimble",
+    "threshing machine",
+    "throne",
+    "tile roof",
+    "toaster",
+    "tobacco shop",
+    "toilet seat",
+    "torch",
+    "totem pole",
+    "tow truck",
+    "toy store",
+    "tractor",
+    "semi-trailer truck",
+    "tray",
+    "trench coat",
+    "tricycle",
+    "trimaran",
+    "tripod",
+    "triumphal arch",
+    "trolleybus",
+    "trombone",
+    "hot tub",
+    "turnstile",
+    "typewriter keyboard",
+    "umbrella",
+    "unicycle",
+    "upright piano",
+    "vacuum cleaner",
+    "vase",
+    "vaulted or arched ceiling",
+    "velvet fabric",
+    "vending machine",
+    "vestment",
+    "viaduct",
+    "violin",
+    "volleyball",
+    "waffle iron",
+    "wall clock",
+    "wallet",
+    "wardrobe",
+    "military aircraft",
+    "sink",
+    "washing machine",
+    "water bottle",
+    "water jug",
+    "water tower",
+    "whiskey jug",
+    "whistle",
+    "hair wig",
+    "window screen",
+    "window shade",
+    "Windsor tie",
+    "wine bottle",
+    "airplane wing",
+    "wok",
+    "wooden spoon",
+    "wool",
+    "split-rail fence",
+    "shipwreck",
+    "sailboat",
+    "yurt",
+    "website",
+    "comic book",
+    "crossword",
+    "traffic or street sign",
+    "traffic light",
+    "dust jacket",
+    "menu",
+    "plate",
+    "guacamole",
+    "consomme",
+    "hot pot",
+    "trifle",
+    "ice cream",
+    "popsicle",
+    "baguette",
+    "bagel",
+    "pretzel",
+    "cheeseburger",
+    "hot dog",
+    "mashed potatoes",
+    "cabbage",
+    "broccoli",
+    "cauliflower",
+    "zucchini",
+    "spaghetti squash",
+    "acorn squash",
+    "butternut squash",
+    "cucumber",
+    "artichoke",
+    "bell pepper",
+    "cardoon",
+    "mushroom",
+    "Granny Smith apple",
+    "strawberry",
+    "orange",
+    "lemon",
+    "fig",
+    "pineapple",
+    "banana",
+    "jackfruit",
+    "cherimoya (custard apple)",
+    "pomegranate",
+    "hay",
+    "carbonara",
+    "chocolate syrup",
+    "dough",
+    "meatloaf",
+    "pizza",
+    "pot pie",
+    "burrito",
+    "red wine",
+    "espresso",
+    "tea cup",
+    "eggnog",
+    "mountain",
+    "bubble",
+    "cliff",
+    "coral reef",
+    "geyser",
+    "lakeshore",
+    "promontory",
+    "sandbar",
+    "beach",
+    "valley",
+    "volcano",
+    "baseball player",
+    "bridegroom",
+    "scuba diver",
+    "rapeseed",
+    "daisy",
+    "yellow lady's slipper",
+    "corn",
+    "acorn",
+    "rose hip",
+    "horse chestnut seed",
+    "coral fungus",
+    "agaric",
+    "gyromitra",
+    "stinkhorn mushroom",
+    "earth star fungus",
+    "hen of the woods mushroom",
+    "bolete",
+    "corn cob",
+    "toilet paper",
+]
+
+HM_CLASSNAMES = [
+    "no",
+    "yes",
+]

open_flamingo/eval/coco_metric.py

@@ -0,0 +1,22 @@
+from pycocoevalcap.eval import COCOEvalCap
+from pycocotools.coco import COCO
+
+
+def compute_cider(
+    result_path,
+    annotations_path,
+):
+    # create coco object and coco_result object
+    coco = COCO(annotations_path)
+    coco_result = coco.loadRes(result_path)
+
+    # create coco_eval object by taking coco and coco_result
+    coco_eval = COCOEvalCap(coco, coco_result)
+    coco_eval.params["image_id"] = coco_result.getImgIds()
+    coco_eval.evaluate()
+
+    return coco_eval.eval
+
+
+def postprocess_captioning_generation(predictions):
+    return predictions.split("Output", 1)[0]

open_flamingo/eval/eval_datasets.py

@@ -0,0 +1,157 @@
+import json
+import os
+
+from PIL import Image
+from torch.utils.data import Dataset
+from torchvision.datasets import ImageFolder
+
+from open_flamingo.eval.classification_utils import IMAGENET_CLASSNAMES
+
+
+class CaptionDataset(Dataset):
+    def __init__(
+        self,
+        image_train_dir_path,
+        annotations_path,
+        is_train,
+        dataset_name,
+        image_val_dir_path=None,
+    ):
+        self.image_train_dir_path = image_train_dir_path
+        self.image_val_dir_path = image_val_dir_path
+        self.annotations = []
+        self.is_train = is_train
+        self.dataset_name = dataset_name
+
+        full_annotations = json.load(open(annotations_path))["images"]
+
+        for i in range(len(full_annotations)):
+            if self.is_train and full_annotations[i]["split"] != "train":
+                continue
+            elif not self.is_train and full_annotations[i]["split"] != "test":
+                continue
+
+            self.annotations.append(full_annotations[i])
+
+    def __len__(self):
+        return len(self.annotations)
+
+    def __getitem__(self, idx):
+        if self.dataset_name == "coco":
+            image = Image.open(
+                os.path.join(
+                    self.image_train_dir_path, self.annotations[idx]["filename"]
+                )
+                if self.annotations[idx]["filepath"] == "train2014"
+                else os.path.join(
+                    self.image_val_dir_path, self.annotations[idx]["filename"]
+                )
+            )
+        elif self.dataset_name == "flickr":
+            image = Image.open(
+                os.path.join(
+                    self.image_train_dir_path, self.annotations[idx]["filename"]
+                )
+            )
+        image.load()
+        caption = self.annotations[idx]["sentences"][0]["raw"]
+        return {
+            "image": image,
+            "caption": caption,
+            "image_id": self.annotations[idx]["cocoid"]
+            if self.dataset_name == "coco"
+            else self.annotations[idx]["filename"].split(".")[0],
+        }
+
+
+class VQADataset(Dataset):
+    def __init__(
+        self, image_dir_path, question_path, annotations_path, is_train, dataset_name
+    ):
+        self.questions = json.load(open(question_path, "r"))["questions"]
+        if annotations_path is not None:
+            self.answers = json.load(open(annotations_path, "r"))["annotations"]
+        else:
+            self.answers = None
+        self.image_dir_path = image_dir_path
+        self.is_train = is_train
+        self.dataset_name = dataset_name
+        if self.dataset_name in {"vqav2", "ok_vqa"}:
+            self.img_coco_split = self.image_dir_path.strip("/").split("/")[-1]
+            assert self.img_coco_split in {"train2014", "val2014", "test2015"}
+
+    def __len__(self):
+        return len(self.questions)
+
+    def get_img_path(self, question):
+        if self.dataset_name in {"vqav2", "ok_vqa"}:
+            return os.path.join(
+                self.image_dir_path,
+                f"COCO_{self.img_coco_split}_{question['image_id']:012d}.jpg"
+                if self.is_train
+                else f"COCO_{self.img_coco_split}_{question['image_id']:012d}.jpg",
+            )
+        elif self.dataset_name == "vizwiz":
+            return os.path.join(self.image_dir_path, question["image_id"])
+        elif self.dataset_name == "textvqa":
+            return os.path.join(self.image_dir_path, f"{question['image_id']}.jpg")
+        else:
+            raise Exception(f"Unknown VQA dataset {self.dataset_name}")
+
+    def __getitem__(self, idx):
+        question = self.questions[idx]
+        img_path = self.get_img_path(question)
+        image = Image.open(img_path)
+        image.load()
+        results = {
+            "image": image,
+            "question": question["question"],
+            "question_id": question["question_id"],
+        }
+        if self.answers is not None:
+            answers = self.answers[idx]
+            results["answers"] = [a["answer"] for a in answers["answers"]]
+        return results
+
+
+class ImageNetDataset(ImageFolder):
+    """Class to represent the ImageNet1k dataset."""
+
+    def __init__(self, root, **kwargs):
+        super().__init__(root=root, **kwargs)
+        self.class_id_to_name = dict(
+            zip(range(len(IMAGENET_CLASSNAMES)), IMAGENET_CLASSNAMES)
+        )
+
+    def __getitem__(self, idx):
+        sample, target = super().__getitem__(idx)
+        target_label = self.class_id_to_name[target]
+        return {
+            "id": idx,
+            "image": sample,
+            "class_id": target,  # numeric ID of the ImageNet class
+            "class_name": target_label,  # human-readable name of ImageNet class
+        }
+
+
+class HatefulMemesDataset(Dataset):
+    def __init__(self, image_dir_path, annotations_path):
+        self.image_dir_path = image_dir_path
+        with open(annotations_path, "r") as f:
+            self.annotations = [json.loads(line) for line in f]
+
+    def __len__(self):
+        return len(self.annotations)
+
+    def __getitem__(self, idx):
+        annotation = self.annotations[idx]
+        img_path = os.path.join(self.image_dir_path, annotation["img"].split("/")[-1])
+        image = Image.open(img_path)
+        image.load()
+        return {
+            "id": annotation["id"],
+            "image": image,
+            "ocr": annotation["text"],
+            "class_name": "yes" if annotation["label"] == 1 else "no",
+            "class_id": annotation["label"],
+        }

open_flamingo/eval/eval_model.py

@@ -0,0 +1,89 @@
+import abc
+import argparse
+from typing import List
+from torch.nn.parallel import DistributedDataParallel as DDP
+from PIL import Image
+
+
+class BaseEvalModel(abc.ABC):
+    """Base class encapsulating functionality needed to evaluate a model."""
+
+    def __init__(self, args: List[str]):
+        """Initialize model.
+
+        Args:
+            args: arguments to model. These should be parsed, or if the model
+                has no applicable arguments, an error should be thrown if `args`
+                is non-empty.
+        """
+
+    def init_distributed(self):
+        """Wrap model as DDP."""
+        self.model = DDP(self.model, device_ids=[self.device])
+
+    def set_device(self, device):
+        """Set device for model."""
+        self.device = device
+        self.model = self.model.to(device)
+
+    def get_outputs(
+        self,
+        batch_text: List[str],
+        batch_images: List[List[Image.Image]],
+        min_generation_length: int,
+        max_generation_length: int,
+        num_beams: int,
+        length_penalty: float,
+    ) -> List[str]:
+        """Get outputs for a batch of images and text.
+
+        Args:
+            batch_text: list of text strings, with the text "<image>" in place
+                of any images to be included.
+            batch_images: images to provide to model. Should be a list of lists,
+              where each list contains the images for a single example.
+            max_generation_length: maximum length of the generated caption.
+                Defaults to 10.
+            num_beams: number of beams to use for beam search. Defaults to 3.
+            length_penalty: length penalty for beam search. Defaults to -2.0.
+
+        Returns:
+            List of decoded output strings.
+        """
+
+    def vqa_prompt(self, question, answer=None) -> str:
+        """Get the prompt to use for VQA evaluation. If the answer is not provided, it should be left blank to be generated by the model.
+
+        Returns:
+            The prompt to use for VQA.
+        """
+
+    def caption_prompt(self, caption=None) -> str:
+        """Get the prompt to use for caption evaluation. If the caption is not provided, it should be left blank to be generated by the model.
+
+        Returns:
+            The prompt to use for captioning.
+        """
+
+    def get_rank_classifications(
+        self,
+        batch_text: List[str],
+        batch_images: List[List[Image.Image]],
+        all_class_names: List[str],
+        use_cache: bool,
+        normalize_length: bool,
+    ):
+        """
+        Returns a (B, |all_class_names|) tensor containing the logprobs for each class name.
+        Args:
+            batch_text: list of text strings, with the text "<image>" in place
+                of any images to be included.
+            batch_images: images to provide to model. Should be a list of lists,
+                where each list contains the images for a single example.
+            all_class_names: list of all class names.
+            use_cache: whether to cache the context to speed up evaluations.
+            normalize_length: whether to normalize logprobs by the length of the
+                class name
+        Returns:
+            (B, |all_class_names|) tensor containing the logprobs for each class name.
+        """

open_flamingo/eval/evaluate.py

@@ -0,0 +1,1301 @@
+import argparse
+import importlib
+import json
+import os
+import uuid
+import random
+from collections import defaultdict
+
+import numpy as np
+import torch
+from sklearn.metrics import roc_auc_score
+import utils
+import math
+
+from coco_metric import compute_cider, postprocess_captioning_generation
+from eval_datasets import (
+    CaptionDataset,
+    VQADataset,
+    ImageNetDataset,
+    HatefulMemesDataset,
+)
+from rices import RICES
+from tqdm import tqdm
+
+
+from classification_utils import (
+    IMAGENET_CLASSNAMES,
+    HM_CLASSNAMES,
+)
+
+from eval_model import BaseEvalModel
+
+from ok_vqa_utils import postprocess_ok_vqa_generation
+from open_flamingo.src.flamingo import Flamingo
+from vqa_metric import compute_vqa_accuracy, postprocess_vqa_generation
+
+from open_flamingo.train.distributed import init_distributed_device, world_info_from_env
+
+parser = argparse.ArgumentParser()
+
+parser.add_argument(
+    "--model",
+    type=str,
+    help="Model name. Currently only `OpenFlamingo` is supported.",
+    default="open_flamingo",
+)
+parser.add_argument(
+    "--results_file", type=str, default=None, help="JSON file to save results"
+)
+
+# Trial arguments
+parser.add_argument("--shots", nargs="+", default=[0, 4, 8, 16, 32], type=int)
+parser.add_argument(
+    "--num_trials",
+    type=int,
+    default=1,
+    help="Number of trials to run for each shot using different demonstrations",
+)
+parser.add_argument(
+    "--trial_seeds",
+    nargs="+",
+    type=int,
+    default=[42],
+    help="Seeds to use for each trial for picking demonstrations and eval sets",
+)
+parser.add_argument(
+    "--num_samples",
+    type=int,
+    default=-1,
+    help="Number of samples to evaluate on. -1 for all samples.",
+)
+parser.add_argument(
+    "--query_set_size", type=int, default=2048, help="Size of demonstration query set"
+)
+
+parser.add_argument("--batch_size", type=int, default=8)
+
+parser.add_argument(
+    "--no_caching_for_classification",
+    action="store_true",
+    help="Whether to skip using key-value caching for classification evals, which usually speeds it up.",
+)
+parser.add_argument(
+    "--classification_prompt_ensembling",
+    action="store_true",
+    help="Whether to use prompt ensembling (average log-likelihoods over permutations of in-context examples)",
+)
+parser.add_argument(
+    "--rices",
+    action="store_true",
+    help="Whether to use RICES for evaluation. If False, uses random demonstrations.",
+)
+parser.add_argument(
+    "--rices_vision_encoder_path",
+    default="ViT-L-14",
+    type=str,
+    help="CLIP vision encoder to use for RICES if cached_demonstration_features is None.",
+)
+parser.add_argument(
+    "--rices_vision_encoder_pretrained",
+    default="openai",
+    type=str,
+    help="CLIP vision encoder to use for RICES if cached_demonstration_features is None.",
+)
+parser.add_argument(
+    "--cached_demonstration_features",
+    default=None,
+    help="Directory where rices features for all choices of in-context examples are stored as a pkl file with the dataset name. If None, features are re-computed by script.",
+)
+
+# Per-dataset evaluation flags
+parser.add_argument(
+    "--eval_coco",
+    action="store_true",
+    default=False,
+    help="Whether to evaluate on COCO.",
+)
+parser.add_argument(
+    "--eval_vqav2",
+    action="store_true",
+    default=False,
+    help="Whether to evaluate on VQAV2.",
+)
+parser.add_argument(
+    "--eval_ok_vqa",
+    action="store_true",
+    default=False,
+    help="Whether to evaluate on OK-VQA.",
+)
+parser.add_argument(
+    "--eval_vizwiz",
+    action="store_true",
+    default=False,
+    help="Whether to evaluate on VizWiz.",
+)
+parser.add_argument(
+    "--eval_textvqa",
+    action="store_true",
+    default=False,
+    help="Whether to evaluate on TextVQA.",
+)
+parser.add_argument(
+    "--eval_imagenet",
+    action="store_true",
+    default=False,
+    help="Whether to evaluate on ImageNet.",
+)
+parser.add_argument(
+    "--eval_flickr30",
+    action="store_true",
+    default=False,
+    help="Whether to evaluate on Flickr30.",
+)
+parser.add_argument(
+    "--eval_hateful_memes",
+    action="store_true",
+    default=False,
+    help="Whether to evaluate on Hateful Memes.",
+)
+
+# Dataset arguments
+
+## Flickr30 Dataset
+parser.add_argument(
+    "--flickr_image_dir_path",
+    type=str,
+    help="Path to the flickr30/flickr30k_images directory.",
+    default=None,
+)
+parser.add_argument(
+    "--flickr_karpathy_json_path",
+    type=str,
+    help="Path to the dataset_flickr30k.json file.",
+    default=None,
+)
+parser.add_argument(
+    "--flickr_annotations_json_path",
+    type=str,
+    help="Path to the dataset_flickr30k_coco_style.json file.",
+)
+## COCO Dataset
+parser.add_argument(
+    "--coco_train_image_dir_path",
+    type=str,
+    default=None,
+)
+parser.add_argument(
+    "--coco_val_image_dir_path",
+    type=str,
+    default=None,
+)
+parser.add_argument(
+    "--coco_karpathy_json_path",
+    type=str,
+    default=None,
+)
+parser.add_argument(
+    "--coco_annotations_json_path",
+    type=str,
+    default=None,
+)
+
+## VQAV2 Dataset
+parser.add_argument(
+    "--vqav2_train_image_dir_path",
+    type=str,
+    default=None,
+)
+parser.add_argument(
+    "--vqav2_train_questions_json_path",
+    type=str,
+    default=None,
+)
+parser.add_argument(
+    "--vqav2_train_annotations_json_path",
+    type=str,
+    default=None,
+)
+parser.add_argument(
+    "--vqav2_test_image_dir_path",
+    type=str,
+    default=None,
+)
+parser.add_argument(
+    "--vqav2_test_questions_json_path",
+    type=str,
+    default=None,
+)
+parser.add_argument(
+    "--vqav2_test_annotations_json_path",
+    type=str,
+    default=None,
+)
+parser.add_argument(
+    "--vqav2_final_test_questions_json_path",
+    type=str,
+    help="Path to the v2_OpenEnded_mscoco_test2015_questions.json file containing all test questions. This is required to format the predictions for EvalAI.",
+    default=None,
+)
+
+## OK-VQA Dataset
+parser.add_argument(
+    "--ok_vqa_train_image_dir_path",
+    type=str,
+    help="Path to the vqav2/train2014 directory.",
+    default=None,
+)
+parser.add_argument(
+    "--ok_vqa_train_questions_json_path",
+    type=str,
+    help="Path to the v2_OpenEnded_mscoco_train2014_questions.json file.",
+    default=None,
+)
+parser.add_argument(
+    "--ok_vqa_train_annotations_json_path",
+    type=str,
+    help="Path to the v2_mscoco_train2014_annotations.json file.",
+    default=None,
+)
+parser.add_argument(
+    "--ok_vqa_test_image_dir_path",
+    type=str,
+    help="Path to the vqav2/val2014 directory.",
+    default=None,
+)
+parser.add_argument(
+    "--ok_vqa_test_questions_json_path",
+    type=str,
+    help="Path to the v2_OpenEnded_mscoco_val2014_questions.json file.",
+    default=None,
+)
+parser.add_argument(
+    "--ok_vqa_test_annotations_json_path",
+    type=str,
+    help="Path to the v2_mscoco_val2014_annotations.json file.",
+    default=None,
+)
+
+## VizWiz Dataset
+parser.add_argument(
+    "--vizwiz_train_image_dir_path",
+    type=str,
+    help="Path to the vizwiz train images directory.",
+    default=None,
+)
+parser.add_argument(
+    "--vizwiz_test_image_dir_path",
+    type=str,
+    help="Path to the vizwiz test images directory.",
+    default=None,
+)
+parser.add_argument(
+    "--vizwiz_train_questions_json_path",
+    type=str,
+    help="Path to the vizwiz questions json file.",
+    default=None,
+)
+parser.add_argument(
+    "--vizwiz_train_annotations_json_path",
+    type=str,
+    help="Path to the vizwiz annotations json file.",
+    default=None,
+)
+parser.add_argument(
+    "--vizwiz_test_questions_json_path",
+    type=str,
+    help="Path to the vizwiz questions json file.",
+    default=None,
+)
+parser.add_argument(
+    "--vizwiz_test_annotations_json_path",
+    type=str,
+    help="Path to the vizwiz annotations json file.",
+    default=None,
+)
+
+# TextVQA Dataset
+parser.add_argument(
+    "--textvqa_image_dir_path",
+    type=str,
+    help="Path to the textvqa images directory.",
+    default=None,
+)
+parser.add_argument(
+    "--textvqa_train_questions_json_path",
+    type=str,
+    help="Path to the textvqa questions json file.",
+    default=None,
+)
+parser.add_argument(
+    "--textvqa_train_annotations_json_path",
+    type=str,
+    help="Path to the textvqa annotations json file.",
+    default=None,
+)
+parser.add_argument(
+    "--textvqa_test_questions_json_path",
+    type=str,
+    help="Path to the textvqa questions json file.",
+    default=None,
+)
+parser.add_argument(
+    "--textvqa_test_annotations_json_path",
+    type=str,
+    help="Path to the textvqa annotations json file.",
+    default=None,
+)
+
+## Imagenet dataset
+parser.add_argument("--imagenet_root", type=str, default="/tmp")
+
+## Hateful Memes dataset
+parser.add_argument(
+    "--hateful_memes_image_dir_path",
+    type=str,
+    default=None,
+)
+parser.add_argument(
+    "--hateful_memes_train_annotations_json_path",
+    type=str,
+    default=None,
+)
+parser.add_argument(
+    "--hateful_memes_test_annotations_json_path",
+    type=str,
+    default=None,
+)
+
+# Distributed evaluation
+parser.add_argument(
+    "--dist-url",
+    default="env://",
+    type=str,
+    help="url used to set up distributed training",
+)
+parser.add_argument(
+    "--dist-backend", default="nccl", type=str, help="distributed backend"
+)
+parser.add_argument(
+    "--horovod",
+    default=False,
+    action="store_true",
+    help="Use horovod for distributed training.",
+)
+parser.add_argument(
+    "--no-set-device-rank",
+    default=False,
+    action="store_true",
+    help="Don't set device index from local rank (when CUDA_VISIBLE_DEVICES restricted to one per proc).",
+)
+
+
+def main():
+    args, leftovers = parser.parse_known_args()
+    module = importlib.import_module(f"open_flamingo.eval.models.{args.model}")
+
+    model_args = {
+        leftovers[i].lstrip("-"): leftovers[i + 1] for i in range(0, len(leftovers), 2)
+    }
+    eval_model = module.EvalModel(model_args)
+
+    # set up distributed evaluation
+    args.local_rank, args.rank, args.world_size = world_info_from_env()
+    device_id = init_distributed_device(args)
+    eval_model.set_device(device_id)
+    eval_model.init_distributed()
+
+    if args.model != "open_flamingo" and args.shots != [0]:
+        raise ValueError("Only 0 shot eval is supported for non-open_flamingo models")
+
+    if len(args.trial_seeds) != args.num_trials:
+        raise ValueError("Number of trial seeds must be == number of trials.")
+
+    results = defaultdict(list)
+
+    if args.eval_flickr30:
+        print("Evaluating on Flickr30k...")
+
+        # load cached demonstration features for RICES
+        if args.cached_demonstration_features is not None:
+            cached_features = torch.load(
+                f"{args.cached_demonstration_features}/flickr30.pkl", map_location="cpu"
+            )
+        else:
+            cached_features = None
+
+        for shot in args.shots:
+            scores = []
+            for seed, trial in zip(args.trial_seeds, range(args.num_trials)):
+                cider_score = evaluate_captioning(
+                    args,
+                    eval_model=eval_model,
+                    num_shots=shot,
+                    seed=seed,
+                    dataset_name="flickr",
+                    cached_features=cached_features,
+                )
+                if args.rank == 0:
+                    print(f"Shots {shot} Trial {trial} CIDEr score: {cider_score}")
+                    scores.append(cider_score)
+
+            if args.rank == 0:
+                print(f"Shots {shot} Mean CIDEr score: {np.nanmean(scores)}")
+                results["flickr30"].append(
+                    {
+                        "shots": shot,
+                        "trials": scores,
+                        "mean": np.nanmean(scores),
+                        "stddev": np.nanstd(scores),
+                    }
+                )
+
+    if args.eval_coco:
+        print("Evaluating on COCO...")
+
+        # load cached demonstration features for RICES
+        if args.cached_demonstration_features is not None:
+            cached_features = torch.load(
+                f"{args.cached_demonstration_features}/coco.pkl", map_location="cpu"
+            )
+        else:
+            cached_features = None
+
+        for shot in args.shots:
+            scores = []
+            for seed, trial in zip(args.trial_seeds, range(args.num_trials)):
+                cider_score = evaluate_captioning(
+                    args,
+                    eval_model=eval_model,
+                    num_shots=shot,
+                    seed=seed,
+                    dataset_name="coco",
+                    cached_features=cached_features,
+                )
+                if args.rank == 0:
+                    print(f"Shots {shot} Trial {trial} CIDEr score: {cider_score}")
+                    scores.append(cider_score)
+
+            if args.rank == 0:
+                print(f"Shots {shot} Mean CIDEr score: {np.nanmean(scores)}")
+                results["coco"].append(
+                    {
+                        "shots": shot,
+                        "trials": scores,
+                        "mean": np.nanmean(scores),
+                        "stddev": np.nanstd(scores),
+                    }
+                )
+
+    if args.eval_ok_vqa:
+        print("Evaluating on OK-VQA...")
+
+        # load cached demonstration features for RICES
+        if args.cached_demonstration_features is not None:
+            cached_features = torch.load(
+                f"{args.cached_demonstration_features}/ok_vqa.pkl", map_location="cpu"
+            )
+        else:
+            cached_features = None
+
+        for shot in args.shots:
+            scores = []
+            for seed, trial in zip(args.trial_seeds, range(args.num_trials)):
+                ok_vqa_score = evaluate_vqa(
+                    args=args,
+                    eval_model=eval_model,
+                    num_shots=shot,
+                    seed=seed,
+                    dataset_name="ok_vqa",
+                    cached_features=cached_features,
+                )
+                if args.rank == 0:
+                    print(f"Shots {shot} Trial {trial} OK-VQA score: {ok_vqa_score}")
+                    scores.append(ok_vqa_score)
+
+            if args.rank == 0:
+                print(f"Shots {shot} Mean OK-VQA score: {np.nanmean(scores)}")
+                results["ok_vqa"].append(
+                    {
+                        "shots": shot,
+                        "trials": scores,
+                        "mean": np.nanmean(scores),
+                        "stddev": np.nanstd(scores),
+                    }
+                )
+
+    if args.eval_vqav2:
+        print("Evaluating on VQAv2...")
+
+        # load cached demonstration features for RICES
+        if args.cached_demonstration_features is not None:
+            cached_features = torch.load(
+                f"{args.cached_demonstration_features}/vqav2.pkl", map_location="cpu"
+            )
+        else:
+            cached_features = None
+
+        for shot in args.shots:
+            scores = []
+            for seed, trial in zip(args.trial_seeds, range(args.num_trials)):
+                vqa_score = evaluate_vqa(
+                    args=args,
+                    eval_model=eval_model,
+                    num_shots=shot,
+                    seed=seed,
+                    dataset_name="vqav2",
+                    cached_features=cached_features,
+                )
+                if args.rank == 0 and vqa_score is not None:
+                    print(f"Shots {shot} Trial {trial} VQA score: {vqa_score}")
+                    scores.append(vqa_score)
+
+            if args.rank == 0 and len(scores) > 0:
+                print(f"Shots {shot} Mean VQA score: {np.nanmean(scores)}")
+                results["vqav2"].append(
+                    {
+                        "shots": shot,
+                        "trials": scores,
+                        "mean": np.nanmean(scores),
+                        "stddev": np.nanstd(scores),
+                    }
+                )
+
+    if args.eval_vizwiz:
+        print("Evaluating on VizWiz...")
+
+        # load cached demonstration features for RICES
+        if args.cached_demonstration_features is not None:
+            cached_features = torch.load(
+                f"{args.cached_demonstration_features}/vizwiz.pkl", map_location="cpu"
+            )
+        else:
+            cached_features = None
+
+        for shot in args.shots:
+            scores = []
+            for seed, trial in zip(args.trial_seeds, range(args.num_trials)):
+                vizwiz_score = evaluate_vqa(
+                    args=args,
+                    eval_model=eval_model,
+                    num_shots=shot,
+                    seed=seed,
+                    dataset_name="vizwiz",
+                    cached_features=cached_features,
+                )
+                if args.rank == 0 and vizwiz_score is not None:
+                    print(f"Shots {shot} Trial {trial} VizWiz score: {vizwiz_score}")
+                    scores.append(vizwiz_score)
+
+            if args.rank == 0 and len(scores) > 0:
+                print(f"Shots {shot} Mean VizWiz score: {np.nanmean(scores)}")
+                results["vizwiz"].append(
+                    {
+                        "shots": shot,
+                        "trials": scores,
+                        "mean": np.nanmean(scores),
+                        "stddev": np.nanstd(scores),
+                    }
+                )
+
+    if args.eval_textvqa:
+        print("Evaluating on TextVQA...")
+
+        # load cached demonstration features for RICES
+        if args.cached_demonstration_features is not None:
+            cached_features = torch.load(
+                f"{args.cached_demonstration_features}/textvqa.pkl", map_location="cpu"
+            )
+        else:
+            cached_features = None
+
+        for shot in args.shots:
+            scores = []
+            for seed, trial in zip(args.trial_seeds, range(args.num_trials)):
+                textvqa_score = evaluate_vqa(
+                    args=args,
+                    eval_model=eval_model,
+                    num_shots=shot,
+                    seed=seed,
+                    dataset_name="textvqa",
+                    max_generation_length=10,
+                    cached_features=cached_features,
+                )
+                if args.rank == 0:
+                    print(f"Shots {shot} Trial {trial} TextVQA score: {textvqa_score}")
+                    scores.append(textvqa_score)
+
+            if args.rank == 0:
+                print(f"Shots {shot} Mean TextVQA score: {np.nanmean(scores)}")
+                results["textvqa"].append(
+                    {
+                        "shots": shot,
+                        "trials": scores,
+                        "mean": np.nanmean(scores),
+                        "stddev": np.nanstd(scores),
+                    }
+                )
+
+    if args.eval_imagenet:
+        print("Evaluating on ImageNet...")
+
+        # load cached demonstration features for RICES
+        if args.cached_demonstration_features is not None:
+            cached_features = torch.load(
+                f"{args.cached_demonstration_features}/imagenet.pkl", map_location="cpu"
+            )
+        else:
+            cached_features = None
+
+        for shot in args.shots:
+            scores = []
+            for seed, trial in zip(args.trial_seeds, range(args.num_trials)):
+                imagenet_score = evaluate_classification(
+                    args,
+                    eval_model=eval_model,
+                    num_shots=shot,
+                    seed=seed,
+                    no_kv_caching=args.no_caching_for_classification,
+                    dataset_name="imagenet",
+                    cached_features=cached_features,
+                    use_prompt_ensembling=args.classification_prompt_ensembling,
+                )
+                if args.rank == 0:
+                    print(
+                        f"Shots {shot} Trial {trial} "
+                        f"ImageNet score: {imagenet_score}"
+                    )
+                    scores.append(imagenet_score)
+
+            if args.rank == 0:
+                print(f"Shots {shot} Mean ImageNet score: {np.nanmean(scores)}")
+                results["imagenet"].append(
+                    {
+                        "shots": shot,
+                        "trials": scores,
+                        "mean": np.nanmean(scores),
+                        "stddev": np.nanstd(scores),
+                    }
+                )
+
+    if args.eval_hateful_memes:
+        print("Evaluating on Hateful Memes...")
+
+        # load cached demonstration features for RICES
+        if args.cached_demonstration_features is not None:
+            cached_features = torch.load(
+                f"{args.cached_demonstration_features}/hateful_memes.pkl",
+                map_location="cpu",
+            )
+        else:
+            cached_features = None
+
+        for shot in args.shots:
+            scores = []
+            for seed, trial in zip(args.trial_seeds, range(args.num_trials)):
+                hateful_memes_score = evaluate_classification(
+                    args,
+                    eval_model=eval_model,
+                    num_shots=shot,
+                    seed=seed,
+                    no_kv_caching=args.no_caching_for_classification,
+                    dataset_name="hateful_memes",
+                    cached_features=cached_features,
+                )
+                if args.rank == 0:
+                    print(
+                        f"Shots {shot} Trial {trial} "
+                        f"Hateful Memes score: {hateful_memes_score}"
+                    )
+                    scores.append(hateful_memes_score)
+
+            if args.rank == 0:
+                print(f"Shots {shot} Mean Hateful Memes score: {np.nanmean(scores)}")
+                results["hateful_memes"].append(
+                    {
+                        "shots": shot,
+                        "trials": scores,
+                        "mean": np.nanmean(scores),
+                        "stddev": np.nanstd(scores),
+                    }
+                )
+
+    if args.rank == 0 and args.results_file is not None:
+        with open(args.results_file, "w") as f:
+            json.dump(results, f)
+
+
+def evaluate_captioning(
+    args: argparse.Namespace,
+    eval_model: BaseEvalModel,
+    seed: int = 42,
+    min_generation_length: int = 0,
+    max_generation_length: int = 20,
+    num_beams: int = 3,
+    length_penalty: float = 0.0,
+    num_shots: int = 8,
+    dataset_name: str = "coco",
+    cached_features=None,
+):
+    """Evaluate a model on COCO dataset.
+
+    Args:
+        args (argparse.Namespace): arguments
+        eval_model (BaseEvalModel): model to evaluate
+        seed (int, optional): seed for random number generator. Defaults to 42.
+        max_generation_length (int, optional): maximum length of the generated caption. Defaults to 20.
+        num_beams (int, optional): number of beams to use for beam search. Defaults to 3.
+        length_penalty (float, optional): length penalty for beam search. Defaults to -2.0.
+        num_shots (int, optional): number of in-context samples to use. Defaults to 8.
+        dataset_name (str, optional): dataset to evaluate on. Can be "coco" or "flickr". Defaults to "coco".
+        cached_features (tensor, optional): cached demonstration features for RICES. Defaults to None.
+    Returns:
+        float: CIDEr score
+
+    """
+
+    if dataset_name == "coco":
+        image_train_dir_path = args.coco_train_image_dir_path
+        image_val_dir_path = args.coco_val_image_dir_path
+        annotations_path = args.coco_karpathy_json_path
+    elif dataset_name == "flickr":
+        image_train_dir_path = (
+            args.flickr_image_dir_path
+        )  # Note: calling this "train" for consistency with COCO but Flickr only has one split for images
+        image_val_dir_path = None
+        annotations_path = args.flickr_karpathy_json_path
+    else:
+        raise ValueError(f"Unsupported dataset: {dataset_name}")
+
+    train_dataset = CaptionDataset(
+        image_train_dir_path=image_train_dir_path,
+        image_val_dir_path=image_val_dir_path,
+        annotations_path=annotations_path,
+        is_train=True,
+        dataset_name=dataset_name if dataset_name != "nocaps" else "coco",
+    )
+
+    test_dataset = CaptionDataset(
+        image_train_dir_path=image_train_dir_path,
+        image_val_dir_path=image_val_dir_path,
+        annotations_path=annotations_path,
+        is_train=False,
+        dataset_name=dataset_name,
+    )
+
+    effective_num_shots = utils.compute_effective_num_shots(num_shots, args.model)
+
+    np.random.seed(seed)
+    test_dataloader = utils.prepare_eval_samples(
+        test_dataset,
+        args.num_samples if args.num_samples > 0 else len(test_dataset),
+        args.batch_size,
+    )
+
+    if args.rices:
+        rices_dataset = RICES(
+            train_dataset,
+            eval_model.device,
+            args.batch_size,
+            cached_features=cached_features,
+            vision_encoder_path=args.rices_vision_encoder_path,
+            vision_encoder_pretrained=args.rices_vision_encoder_pretrained,
+        )
+    else:
+        # subset of the training set to sample context images from
+        query_set = utils.get_query_set(train_dataset, args.query_set_size)
+
+    utils.random_seed(seed, args.rank)
+    predictions = defaultdict()
+    for batch in tqdm(
+        test_dataloader,
+        desc=f"Running inference {dataset_name.upper()}",
+        disable=args.rank != 0,
+    ):
+        if args.rices:
+            batch_demo_samples = rices_dataset.find(batch["image"], effective_num_shots)
+        else:
+            batch_demo_samples = utils.sample_batch_demos_from_query_set(
+                query_set, effective_num_shots, len(batch["image"])
+            )
+
+        batch_images, batch_text = [], []
+        for i in range(len(batch["image"])):
+            if num_shots > 0:
+                context_images = [x["image"] for x in batch_demo_samples[i]]
+            else:
+                context_images = []
+            batch_images.append(context_images + [batch["image"][i]])
+
+            context_text = "".join(
+                [
+                    eval_model.get_caption_prompt(caption=x["caption"].strip()) + "\n"
+                    for x in batch_demo_samples[i]
+                ]
+            )
+
+            # Keep the text but remove the image tags for the zero-shot case
+            if num_shots == 0:
+                context_text = context_text.replace("<image>", "")
+
+            batch_text.append(context_text + eval_model.get_caption_prompt())
+
+        outputs = eval_model.get_outputs(
+            batch_images=batch_images,
+            batch_text=batch_text,
+            min_generation_length=min_generation_length,
+            max_generation_length=max_generation_length,
+            num_beams=num_beams,
+            length_penalty=length_penalty,
+        )
+
+        new_predictions = [
+            postprocess_captioning_generation(out).replace('"', "") for out in outputs
+        ]
+
+        for i, sample_id in enumerate(batch["image_id"]):
+            predictions[sample_id] = {
+                "caption": new_predictions[i],
+            }
+
+    # all gather
+    all_predictions = [None for _ in range(args.world_size)]
+    torch.distributed.all_gather_object(all_predictions, predictions)  # list of dicts
+
+    if args.rank != 0:
+        return None
+
+    all_predictions = {
+        k: v for d in all_predictions for k, v in d.items()
+    }  # merge dicts
+
+    # save the predictions to a temporary file
+    results_path = f"{dataset_name}results_{uuid.uuid4()}.json"
+
+    with open(results_path, "w") as f:
+        f.write(
+            json.dumps(
+                [
+                    {"image_id": k, "caption": all_predictions[k]["caption"]}
+                    for k in all_predictions
+                ],
+                indent=4,
+            )
+        )
+
+    metrics = compute_cider(
+        result_path=results_path,
+        annotations_path=args.coco_annotations_json_path
+        if dataset_name == "coco"
+        else args.flickr_annotations_json_path,
+    )
+
+    # delete the temporary file
+    os.remove(results_path)
+
+    return metrics["CIDEr"] * 100.0
+
+
+def evaluate_vqa(
+    args: argparse.Namespace,
+    eval_model: BaseEvalModel,
+    seed: int = 42,
+    min_generation_length: int = 0,
+    max_generation_length: int = 5,
+    num_beams: int = 3,
+    length_penalty: float = 0.0,
+    num_shots: int = 8,
+    dataset_name: str = "vqav2",
+    cached_features=None,
+):
+    """
+    Evaluate a model on VQA datasets. Currently supports VQA v2.0, OK-VQA, VizWiz and TextVQA.
+
+    Args:
+        args (argparse.Namespace): arguments
+        eval_model (BaseEvalModel): model to evaluate
+        seed (int, optional): random seed. Defaults to 42.
+        max_generation_length (int, optional): max generation length. Defaults to 5.
+        num_beams (int, optional): number of beams to use for beam search. Defaults to 3.
+        length_penalty (float, optional): length penalty for beam search. Defaults to -2.0.
+        num_shots (int, optional): number of shots to use. Defaults to 8.
+        dataset_name (string): type of vqa dataset: currently supports vqav2, ok_vqa. Defaults to vqav2.
+        cached_features (tensor, optional): cached demonstration features for RICES. Defaults to None.
+    Returns:
+        float: accuracy score
+    """
+
+    if dataset_name == "ok_vqa":
+        train_image_dir_path = args.ok_vqa_train_image_dir_path
+        train_questions_json_path = args.ok_vqa_train_questions_json_path
+        train_annotations_json_path = args.ok_vqa_train_annotations_json_path
+        test_image_dir_path = args.ok_vqa_test_image_dir_path
+        test_questions_json_path = args.ok_vqa_test_questions_json_path
+        test_annotations_json_path = args.ok_vqa_test_annotations_json_path
+    elif dataset_name == "vqav2":
+        train_image_dir_path = args.vqav2_train_image_dir_path
+        train_questions_json_path = args.vqav2_train_questions_json_path
+        train_annotations_json_path = args.vqav2_train_annotations_json_path
+        test_image_dir_path = args.vqav2_test_image_dir_path
+        test_questions_json_path = args.vqav2_test_questions_json_path
+        test_annotations_json_path = args.vqav2_test_annotations_json_path
+    elif dataset_name == "vizwiz":
+        train_image_dir_path = args.vizwiz_train_image_dir_path
+        train_questions_json_path = args.vizwiz_train_questions_json_path
+        train_annotations_json_path = args.vizwiz_train_annotations_json_path
+        test_image_dir_path = args.vizwiz_test_image_dir_path
+        test_questions_json_path = args.vizwiz_test_questions_json_path
+        test_annotations_json_path = args.vizwiz_test_annotations_json_path
+    elif dataset_name == "textvqa":
+        train_image_dir_path = args.textvqa_image_dir_path
+        train_questions_json_path = args.textvqa_train_questions_json_path
+        train_annotations_json_path = args.textvqa_train_annotations_json_path
+        test_image_dir_path = args.textvqa_image_dir_path
+        test_questions_json_path = args.textvqa_test_questions_json_path
+        test_annotations_json_path = args.textvqa_test_annotations_json_path
+    else:
+        raise ValueError(f"Unsupported dataset: {dataset_name}")
+
+    train_dataset = VQADataset(
+        image_dir_path=train_image_dir_path,
+        question_path=train_questions_json_path,
+        annotations_path=train_annotations_json_path,
+        is_train=True,
+        dataset_name=dataset_name,
+    )
+
+    test_dataset = VQADataset(
+        image_dir_path=test_image_dir_path,
+        question_path=test_questions_json_path,
+        annotations_path=test_annotations_json_path,
+        is_train=False,
+        dataset_name=dataset_name,
+    )
+
+    effective_num_shots = utils.compute_effective_num_shots(num_shots, args.model)
+
+    np.random.seed(seed)
+    test_dataloader = utils.prepare_eval_samples(
+        test_dataset,
+        args.num_samples if args.num_samples > 0 else len(test_dataset),
+        args.batch_size,
+    )
+
+    if args.rices:
+        rices_dataset = RICES(
+            train_dataset,
+            eval_model.device,
+            args.batch_size,
+            cached_features=cached_features,
+            vision_encoder_path=args.rices_vision_encoder_path,
+            vision_encoder_pretrained=args.rices_vision_encoder_pretrained,
+        )
+    else:
+        query_set = utils.get_query_set(train_dataset, args.query_set_size)
+
+    utils.random_seed(seed, args.rank)
+    predictions = []
+    for batch in tqdm(
+        test_dataloader,
+        desc=f"Running inference {dataset_name}",
+        disable=args.rank != 0,
+    ):
+        if args.rices:
+            batch_demo_samples = rices_dataset.find(batch["image"], effective_num_shots)
+        else:
+            batch_demo_samples = utils.sample_batch_demos_from_query_set(
+                query_set, effective_num_shots, len(batch["image"])
+            )
+
+        batch_images, batch_text = [], []
+        for i in range(len(batch["image"])):
+            if num_shots > 0:
+                context_images = [x["image"] for x in batch_demo_samples[i]]
+            else:
+                context_images = []
+            batch_images.append(context_images + [batch["image"][i]])
+
+            context_text = "".join(
+                [
+                    eval_model.get_vqa_prompt(
+                        question=x["question"], answer=x["answers"][0]
+                    )
+                    + "\n"
+                    for x in batch_demo_samples[i]
+                ]
+            )
+
+            # Keep the text but remove the image tags for the zero-shot case
+            if num_shots == 0:
+                context_text = context_text.replace("<image>", "")
+
+            batch_text.append(
+                context_text + eval_model.get_vqa_prompt(question=batch["question"][i])
+            )
+
+        outputs = eval_model.get_outputs(
+            batch_images=batch_images,
+            batch_text=batch_text,
+            min_generation_length=min_generation_length,
+            max_generation_length=max_generation_length,
+            num_beams=num_beams,
+            length_penalty=length_penalty,
+        )
+
+        process_function = (
+            postprocess_ok_vqa_generation
+            if dataset_name == "ok_vqa"
+            else postprocess_vqa_generation
+        )
+
+        new_predictions = map(process_function, outputs)
+
+        for new_prediction, sample_id in zip(new_predictions, batch["question_id"]):
+            predictions.append({"answer": new_prediction, "question_id": sample_id})
+
+    # all gather
+    all_predictions = [None for _ in range(args.world_size)]
+    torch.distributed.all_gather_object(all_predictions, predictions)  # list of lists
+
+    if args.rank != 0:
+        return None
+
+    all_predictions = [
+        item for sublist in all_predictions for item in sublist
+    ]  # flatten
+
+    # save the predictions to a temporary file
+    random_uuid = str(uuid.uuid4())
+    with open(f"{dataset_name}results_{random_uuid}.json", "w") as f:
+        f.write(json.dumps(all_predictions, indent=4))
+
+    if test_annotations_json_path is not None:
+        acc = compute_vqa_accuracy(
+            f"{dataset_name}results_{random_uuid}.json",
+            test_questions_json_path,
+            test_annotations_json_path,
+        )
+        # delete the temporary file
+        os.remove(f"{dataset_name}results_{random_uuid}.json")
+
+    else:
+        print("No annotations provided, skipping accuracy computation.")
+        acc = None
+        if dataset_name == "vqav2":
+            from open_flamingo.scripts.fill_vqa_testdev_results import (
+                fill_vqav2_test_json,
+            )
+
+            fill_fn = fill_vqav2_test_json
+        elif dataset_name == "vizwiz":
+            from open_flamingo.scripts.fill_vqa_testdev_results import (
+                fill_vizwiz_test_json,
+            )
+
+            fill_fn = fill_vizwiz_test_json
+        else:
+            print(
+                "Temporary file saved to ", f"{dataset_name}results_{random_uuid}.json"
+            )
+            return
+
+        fill_fn(
+            f"{dataset_name}results_{random_uuid}.json",
+            f"{dataset_name}-testdev_{eval_model.lm_name}_{num_shots}_{'rices' if args.rices else 'random'}_{seed}.json",
+            args.vqav2_final_test_questions_json_path
+            if dataset_name == "vqav2"
+            else args.vizwiz_test_questions_json_path,
+        )
+        print(
+            "Test-dev results saved to ",
+            f"{dataset_name}-testdev_{eval_model.lm_name}_{num_shots}_{'rices' if args.rices else 'random'}_{seed}.json",
+        )
+        os.remove(f"{dataset_name}results_{random_uuid}.json")
+
+    return acc
+
+
+def evaluate_classification(
+    args: argparse.Namespace,
+    eval_model,
+    seed: int = 42,
+    num_shots: int = 8,
+    dataset_name: str = "imagenet",
+    cached_features=None,
+    no_kv_caching=False,
+    use_prompt_ensembling: bool = False,
+):
+    """
+    Evaluate a model on classification dataset.
+
+    Args:
+        eval_model (BaseEvalModel): model to evaluate
+        seed (int, optional): random seed. Defaults to 42.
+        num_shots (int, optional): number of shots to use. Defaults to 8.
+        no_kv_caching (bool): whether to disable key-value caching
+        dataset_name (str, optional): dataset name. Defaults to "imagenet".
+        cached_features (tensor, optional): cached demonstration features for RICES. Defaults to None.
+
+    Returns:
+        float: accuracy score
+    """
+    if args.model != "open_flamingo":
+        raise NotImplementedError(
+            "evaluate_classification is currently only supported for OpenFlamingo"
+        )
+
+    if dataset_name == "imagenet":
+        train_dataset = ImageNetDataset(os.path.join(args.imagenet_root, "train"))
+        test_dataset = ImageNetDataset(os.path.join(args.imagenet_root, "val"))
+        prompt_fn = lambda x: eval_model.get_imagenet_prompt(label=x["class_name"])
+        all_class_names = IMAGENET_CLASSNAMES
+        k = 5
+    elif dataset_name == "hateful_memes":
+        train_dataset = HatefulMemesDataset(
+            args.hateful_memes_image_dir_path,
+            args.hateful_memes_train_annotations_json_path,
+        )
+        test_dataset = HatefulMemesDataset(
+            args.hateful_memes_image_dir_path,
+            args.hateful_memes_test_annotations_json_path,
+        )
+        prompt_fn = lambda x: eval_model.get_hateful_memes_prompt(
+            text=x["ocr"], label=x["class_name"]
+        )
+        all_class_names = HM_CLASSNAMES
+        k = 1
+    else:
+        raise ValueError(f"Unsupported dataset {dataset_name}")
+
+    class_id_to_name = dict(zip(range(len(all_class_names)), all_class_names))
+
+    effective_num_shots = utils.compute_effective_num_shots(num_shots, args.model)
+
+    np.random.seed(seed)
+    test_dataloader = utils.prepare_eval_samples(
+        test_dataset,
+        args.num_samples if args.num_samples > 0 else len(test_dataset),
+        args.batch_size,
+    )
+
+    if args.rices:
+        rices_dataset = RICES(
+            train_dataset,
+            eval_model.device,
+            args.batch_size,
+            cached_features=cached_features,
+            vision_encoder_path=args.rices_vision_encoder_path,
+            vision_encoder_pretrained=args.rices_vision_encoder_pretrained,
+        )
+    else:
+        # subset of the training set to sample context images from
+        query_set = utils.get_query_set(train_dataset, args.query_set_size)
+
+    utils.random_seed(seed, args.rank)
+    predictions = []
+    for batch_idx, batch in tqdm(
+        enumerate(test_dataloader),
+        desc=f"Running inference {dataset_name}",
+        disable=args.rank != 0,
+    ):
+        if args.rices:
+            batch_demo_samples = rices_dataset.find(batch["image"], effective_num_shots)
+        else:
+            batch_demo_samples = utils.sample_batch_demos_from_query_set(
+                query_set, effective_num_shots, len(batch["image"])
+            )
+
+        # set up prompt ensembling
+        num_permutations = (
+            min(6, math.factorial(effective_num_shots)) if use_prompt_ensembling else 1
+        )
+        logprobs = []
+        for _ in range(num_permutations):
+            batch_images, batch_text = [], []
+            for i in range(len(batch["image"])):
+                if use_prompt_ensembling:
+                    random.shuffle(batch_demo_samples[i])
+
+                if effective_num_shots > 0:
+                    context_images = [x["image"] for x in batch_demo_samples[i]]
+                else:
+                    context_images = []
+                batch_images.append(context_images + [batch["image"][i]])
+
+                context_text = "".join([prompt_fn(x) for x in batch_demo_samples[i]])
+
+                # Keep the text but remove the image tags for the zero-shot case
+                if num_shots == 0:
+                    context_text = context_text.replace("<image>", "")
+
+                batch_text.append(
+                    context_text
+                    + prompt_fn({"ocr": batch["ocr"][i], "class_name": None})
+                )
+
+            # get predicted class names
+            logprobs.append(
+                eval_model.get_rank_classifications(
+                    batch_text,
+                    batch_images,
+                    all_class_names,
+                    use_cache=(not no_kv_caching),
+                    normalize_length=True,
+                )
+            )
+
+        # ensemble logprobs together
+        logprobs = torch.mean(torch.stack(logprobs, dim=-1), dim=-1)
+
+        predicted_classnames, predicted_logprobs = utils.get_predicted_classnames(
+            logprobs,
+            k,
+            class_id_to_name,
+        )
+
+        # compute accuracy
+        for i, topk in enumerate(predicted_classnames):
+            y_i = batch["class_name"][i]
+            score = torch.exp(
+                predicted_logprobs[i][0] - torch.logsumexp(logprobs[i], dim=0)
+            ).item()
+            predictions.append(
+                {
+                    "id": batch["id"][i],
+                    "gt_label": y_i,
+                    "pred_label": topk[0],
+                    "pred_score": score,
+                }
+            )
+
+    # all gather
+    all_predictions = [None for _ in range(args.world_size)]
+    torch.distributed.all_gather_object(all_predictions, predictions)  # list of lists
+    if args.rank != 0:
+        return
+
+    all_predictions = [
+        item for sublist in all_predictions for item in sublist
+    ]  # flatten
+
+    if dataset_name == "hateful_memes":
+        # return ROC-AUC score
+        greater_label = max(all_class_names)
+        gts = [pred["gt_label"] for pred in all_predictions]
+        pred_scores = [
+            pred["pred_score"]
+            if pred["pred_label"] == greater_label
+            else 1 - pred["pred_score"]
+            for pred in all_predictions
+        ]
+        return roc_auc_score(gts, pred_scores)
+    else:
+        # return top-1 accuracy
+        acc1 = sum(
+            int(pred["gt_label"] == pred["pred_label"]) for pred in all_predictions
+        )
+        return float(acc1) / len(all_predictions)
+
+
+if __name__ == "__main__":
+    main()

open_flamingo/eval/models/blip.py

@@ -0,0 +1,117 @@
+from typing import List
+
+from PIL import Image
+import torch
+
+from transformers import Blip2Processor, Blip2ForConditionalGeneration
+from open_flamingo.eval.eval_model import BaseEvalModel
+from open_flamingo.eval.utils import unwrap_model
+
+
+class EvalModel(BaseEvalModel):
+    """BLIP-2 model evaluation.
+
+    Attributes:
+      model (nn.Module): Underlying Torch model.
+      tokenizer (transformers.PreTrainedTokenizer): Tokenizer for model.
+      device: Index of GPU to use, or the string "cpu"
+    """
+
+    def __init__(self, model_args):
+        assert (
+            "processor_path" in model_args and "lm_path" in model_args
+        ), "BLIP-2 requires processor_path, lm_path, and device arguments to be specified"
+
+        self.processor = Blip2Processor.from_pretrained(model_args["processor_path"])
+        self.model = Blip2ForConditionalGeneration.from_pretrained(
+            model_args["lm_path"]
+        )
+        self.model.eval()
+        self.processor.tokenizer.padding_side = "left"
+        self.lm_name = model_args["lm_path"].split("/")[-1]
+
+    def _prepare_images(self, batch: List[List[torch.Tensor]]) -> torch.Tensor:
+        """Preprocess images and stack them.
+
+        Args:
+            batch: A list of lists of images.
+
+        Returns:
+            A Tensor of shape
+            (batch_size, channels, height, width).
+        """
+        batch_images = None
+        assert all(
+            len(example) == 1 for example in batch
+        ), "BLIP-2 only supports one image per example"
+
+        for example in batch:
+            assert len(example) == 1, "BLIP-2 only supports one image per example"
+            batch_images = torch.cat(
+                [
+                    batch_images,
+                    self.processor.image_processor(example, return_tensors="pt")[
+                        "pixel_values"
+                    ],
+                ]
+                if batch_images is not None
+                else [
+                    self.processor.image_processor(example, return_tensors="pt")[
+                        "pixel_values"
+                    ]
+                ],
+                dim=0,
+            )
+        return batch_images
+
+    def get_outputs(
+        self,
+        batch_text: List[str],
+        batch_images: List[List[Image.Image]],
+        min_generation_length: int,
+        max_generation_length: int,
+        num_beams: int,
+        length_penalty: float,
+    ) -> List[str]:
+        encodings = self.processor.tokenizer(
+            batch_text,
+            padding="longest",
+            truncation=True,
+            return_tensors="pt",
+            max_length=2000,
+        )
+        input_ids = encodings["input_ids"]
+        attention_mask = encodings["attention_mask"]
+
+        with torch.inference_mode():
+            outputs = unwrap_model(self.model).generate(
+                self._prepare_images(batch_images).to(self.device),
+                input_ids.to(self.device),
+                attention_mask=attention_mask.to(self.device),
+                max_new_tokens=max_generation_length,
+                min_new_tokens=min_generation_length,
+                num_beams=num_beams,
+                length_penalty=length_penalty,
+            )
+
+        return self.processor.tokenizer.batch_decode(outputs, skip_special_tokens=True)
+
+    def get_vqa_prompt(self, question, answer=None) -> str:
+        return (
+            f"Question:{question} Short answer:{answer if answer is not None else ''}"
+        )
+
+    def get_caption_prompt(self, caption=None) -> str:
+        return f"A photo of {caption if caption is not None else ''}"
+
+    def get_rank_classifications(
+        self,
+        batch_text: List[str],
+        batch_images: List[List[Image.Image]],
+        all_class_names: List[str],
+        use_cache: bool,
+        normalize_length: bool,
+    ):
+        raise NotImplementedError(
+            "BLIP-2 classification-based evaluation not implemented"
+        )

open_flamingo/eval/models/open_flamingo.py

@@ -0,0 +1,334 @@
+from typing import List, Dict
+
+from PIL import Image
+import torch
+from einops import repeat
+
+from open_flamingo.eval.eval_model import BaseEvalModel
+from open_flamingo.src.factory import create_model_and_transforms
+from open_flamingo.eval.utils import unwrap_model, get_autocast, get_cast_dtype
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+
+class EvalModel(BaseEvalModel):
+    """OpenFlamingo model evaluation.
+
+    Attributes:
+      model (nn.Module): Underlying Torch model.
+      tokenizer (transformers.PreTrainedTokenizer): Tokenizer for model.
+      device: Index of GPU to use, or the string "CPU"
+    """
+
+    def __init__(self, model_args):
+        assert (
+            "vision_encoder_path" in model_args
+            and "lm_path" in model_args
+            and "checkpoint_path" in model_args
+            and "lm_tokenizer_path" in model_args
+            and "cross_attn_every_n_layers" in model_args
+            and "vision_encoder_pretrained" in model_args
+            and "precision" in model_args
+        ), "OpenFlamingo requires vision_encoder_path, lm_path, device, checkpoint_path, lm_tokenizer_path, cross_attn_every_n_layers, vision_encoder_pretrained, and precision arguments to be specified"
+
+        self.device = (
+            model_args["device"]
+            if ("device" in model_args and model_args["device"] >= 0)
+            else "cpu"
+        )
+
+        (
+            self.model,
+            self.image_processor,
+            self.tokenizer,
+        ) = create_model_and_transforms(
+            model_args["vision_encoder_path"],
+            model_args["vision_encoder_pretrained"],
+            model_args["lm_path"],
+            model_args["lm_tokenizer_path"],
+            cross_attn_every_n_layers=int(model_args["cross_attn_every_n_layers"]),
+        )
+        checkpoint = torch.load(model_args["checkpoint_path"], map_location=self.device)
+        if "model_state_dict" in checkpoint:
+            checkpoint = checkpoint["model_state_dict"]
+            checkpoint = {k.replace("module.", ""): v for k, v in checkpoint.items()}
+        self.model.load_state_dict(checkpoint, strict=False)
+        self.model.to(self.device)
+        self.model.eval()
+        self.tokenizer.padding_side = "left"
+
+        self.lm_name = model_args["lm_path"].split("/")[-1]
+
+        # autocast
+        self.autocast = get_autocast(model_args["precision"])
+        self.cast_dtype = get_cast_dtype(model_args["precision"])
+
+    def _prepare_images(self, batch: List[List[Image.Image]]) -> torch.Tensor:
+        """
+        Convert images to tensors, reshape them, and stack them.
+        Args:
+            batch: A list of lists of images.
+        Returns:
+            preprocessed images (tensors) or None
+                shape (B, T_img, F, C, H, W)
+                None if no images in batch
+        """
+        images_per_example = max(len(x) for x in batch)
+        batch_images = None
+        for iexample, example in enumerate(batch):
+            for iimage, image in enumerate(example):
+                preprocessed = self.image_processor(image)
+                if batch_images is None:
+                    batch_images = torch.zeros(
+                        (len(batch), images_per_example, 1) + preprocessed.shape,
+                        dtype=preprocessed.dtype,
+                    )
+                batch_images[iexample, iimage, 0] = preprocessed
+        if batch_images is not None:
+            batch_images = batch_images.to(
+                self.device, dtype=self.cast_dtype, non_blocking=True
+            )
+        return batch_images
+
+    def _prepare_text(
+        self,
+        batch: List[List[str]],
+        padding="longest",
+        truncation=True,
+        max_length=2000,
+    ):
+        """
+        Tokenize the text and stack them.
+        Args:
+            batch: A list of lists of strings.
+        Returns:
+            input_ids (tensor)
+                shape (B, T_txt)
+            attention_mask (tensor)
+                shape (B, T_txt)
+        """
+        encodings = self.tokenizer(
+            batch,
+            padding=padding,
+            truncation=truncation,
+            return_tensors="pt",
+            max_length=max_length,
+        )
+        input_ids, attention_mask = encodings["input_ids"], encodings["attention_mask"]
+        input_ids = input_ids.to(self.device, dtype=self.cast_dtype, non_blocking=True)
+        attention_mask = attention_mask.to(
+            self.device, dtype=self.cast_dtype, non_blocking=True
+        )
+        return input_ids, attention_mask.bool()
+
+    def get_outputs(
+        self,
+        batch_text: List[str],
+        batch_images: List[List[Image.Image]],
+        min_generation_length: int,
+        max_generation_length: int,
+        num_beams: int,
+        length_penalty: float,
+    ) -> List[str]:
+        """
+        Get generation outputs.
+        """
+        batch_images = self._prepare_images(batch_images)
+        input_ids, attention_mask = self._prepare_text(batch_text)
+
+        with torch.inference_mode():
+            with self.autocast():
+                outputs = unwrap_model(self.model).generate(
+                    batch_images,
+                    input_ids,
+                    attention_mask,
+                    min_new_tokens=min_generation_length,
+                    max_new_tokens=max_generation_length,
+                    num_beams=num_beams,
+                    length_penalty=length_penalty,
+                )
+
+        # Extract only the new gnerated tokens
+        outputs = outputs[:, len(input_ids[0]) :]
+
+        return self.tokenizer.batch_decode(outputs, skip_special_tokens=True)
+
+    def get_rank_classifications(
+        self,
+        batch_text: List[str],
+        batch_images: List[List[Image.Image]],
+        all_class_names: List[str],
+        use_cache: bool,
+        normalize_length: bool,
+    ):
+        """
+        Returns a (B, |all_class_names|) tensor containing the logprobs for each class name.
+        """
+        batch_images = self._prepare_images(batch_images)
+        ctx_input_ids, ctx_attention_mask = self._prepare_text(batch_text)
+
+        # Cache the context
+        if use_cache:
+            # reserve the last token in the context for the main forward pass
+            self.cache_media(
+                input_ids=ctx_input_ids,
+                vision_x=batch_images,
+            )
+            precomputed = self.__call__(
+                vision_x=None,
+                lang_x=ctx_input_ids,
+                attention_mask=ctx_attention_mask,
+                clear_conditioned_layers=False,
+                use_cache=True,
+            )
+            precomputed_logits = precomputed.logits
+            precomputed_pkvs = precomputed.past_key_values
+        else:
+            precomputed_pkvs = None
+
+        # Loop through class names and get log-likelihoods
+        # Note: if all classnames are one token, this code is redundant, since we could
+        # get all logits after one pass. However, if there are multi-token classnames,
+        # we need to loop through each classname separately.
+        overall_probs = []
+        for class_name in all_class_names:
+            # Tokenize only the class name
+            classname_tokens = self.tokenizer(
+                class_name, add_special_tokens=False, return_tensors="pt"
+            )["input_ids"].to(self.device)
+            assert classname_tokens.ndim == 2
+            classname_tokens = repeat(
+                classname_tokens, "b s -> (repeat b) s", repeat=len(batch_text)
+            )
+            num_tokens_in_classname = classname_tokens.shape[1]
+
+            # Concatenate the class name tokens
+            if not use_cache:
+                _lang_x = torch.cat([ctx_input_ids, classname_tokens], dim=1)
+                _attention_mask = torch.cat(
+                    [
+                        ctx_attention_mask,
+                        torch.ones_like(classname_tokens).bool(),
+                    ],
+                    dim=1,
+                )
+                _vision_x = batch_images
+            else:
+                _lang_x = classname_tokens
+                _attention_mask = None
+                _vision_x = None
+
+            # Call forward to get the logits
+            outputs = self.__call__(
+                vision_x=_vision_x,
+                lang_x=_lang_x,
+                attention_mask=_attention_mask,
+                clear_conditioned_layers=(not use_cache),
+                past_key_values=precomputed_pkvs,
+            )
+
+            # Get the logits of the classname
+            # logits shape is either (B, num_tokens_in_classname, vocab_len) with use_cache
+            # or (B, len(_lang_x), vocab_len) without use_cache
+            # remember that the logits at index t on dim 1 correspond to predictions for the t+1st token
+            logits = outputs.logits
+            if use_cache:
+                logits = torch.cat([precomputed_logits, logits], dim=1)
+
+            logprobs = torch.log_softmax(logits, dim=-1)
+            gen_probs = logprobs[
+                :, -num_tokens_in_classname - 1 : -1, :
+            ]  # (B, num_tokens_in_classname, vocab_len)
+            gen_probs = torch.gather(
+                gen_probs, 2, classname_tokens[:, :, None]
+            ).squeeze(-1)
+
+            # Aggregate over tokens in the classname
+            if normalize_length:
+                class_prob = torch.mean(gen_probs, dim=1)
+            else:
+                class_prob = torch.sum(gen_probs, dim=1)
+            overall_probs.append(class_prob)  # (B, 1)
+
+        self.uncache_media()
+        overall_probs = torch.vstack(overall_probs).T.cpu()  # shape (B, num_classes)
+        return overall_probs
+
+    def __call__(
+        self,
+        lang_x: torch.Tensor,
+        vision_x: torch.Tensor,
+        attention_mask: torch.Tensor,
+        past_key_values: torch.Tensor = None,
+        clear_conditioned_layers: bool = False,
+        use_cache: bool = False,
+    ):
+        """
+        Calls the forward function of the model.
+        Special logic to handle the case if past_key_values is not None:
+            then lang_x is assumed to contain the tokens to be generated
+            *excluding* the tokens already in past_key_values.
+            We then repeatedly call forward, updating the past_key_values.
+        """
+        # standard forward pass
+        if past_key_values is None:
+            with torch.inference_mode():
+                with self.autocast():
+                    outputs = self.model(
+                        vision_x=vision_x,
+                        lang_x=lang_x,
+                        attention_mask=attention_mask,
+                        clear_conditioned_layers=clear_conditioned_layers,
+                        past_key_values=past_key_values,
+                        use_cache=use_cache,
+                    )
+            return outputs
+
+        # loop to handle updating past_key_values
+        logits = []
+        for token_idx in range(lang_x.shape[1]):
+            _lang_x = lang_x[:, token_idx].reshape((-1, 1))
+            if attention_mask is not None:
+                _attention_mask = attention_mask[:, token_idx].reshape((-1, 1))
+            else:
+                _attention_mask = None
+
+            with torch.inference_mode():
+                with self.autocast():
+                    outputs = self.model(
+                        vision_x=vision_x,
+                        lang_x=_lang_x,
+                        attention_mask=_attention_mask,
+                        clear_conditioned_layers=False,
+                        past_key_values=past_key_values,
+                        use_cache=True,
+                    )
+
+            past_key_values = outputs.past_key_values
+            logits.append(outputs.logits)
+
+        logits = torch.cat(logits, dim=1)
+        return CausalLMOutputWithPast(
+            logits=logits,
+            past_key_values=past_key_values,
+        )
+
+    def encode_vision_x(self, image_tensor: torch.Tensor):
+        unwrap_model(self.model)._encode_vision_x(image_tensor.to(self.device))
+
+    def uncache_media(self):
+        unwrap_model(self.model).uncache_media()
+
+    def cache_media(self, input_ids, vision_x):
+        unwrap_model(self.model).cache_media(input_ids=input_ids, vision_x=vision_x)
+
+    def get_vqa_prompt(self, question, answer=None) -> str:
+        return f"<image>Question:{question} Short answer:{answer if answer is not None else ''}{'<|endofchunk|>' if answer is not None else ''}"
+
+    def get_caption_prompt(self, caption=None) -> str:
+        return f"<image>Output:{caption if caption is not None else ''}{'<|endofchunk|>' if caption is not None else ''}"
+
+    def get_imagenet_prompt(self, label=None) -> str:
+        return f"<image>Output:{label if label is not None else ''}{'<|endofchunk|>' if label is not None else ''}"
+
+    def get_hateful_memes_prompt(self, text, label=None) -> str:
+        return f"<image>is an image with: '{text}' written on it. Is it hateful? Answer:{label if label is not None else ''}{'<|endofchunk|>' if label is not None else ''}"

open_flamingo/eval/ok_vqa_utils.py

@@ -0,0 +1,215 @@
+# Those are manual mapping that are not caught by our stemming rules or would
+# would be done incorrectly by our automatic stemming rule. In details,
+# the keys of the _MANUAL_MATCHES dict contains the original word and the value
+# contains the transformation of the word expected by the OKVQA stemming rule.
+# These manual rules were found by checking the `raw_answers` and the `answers`
+# fields of the released OKVQA dataset and checking all things that were not
+# properly mapped by our automatic rules. In particular some of the mapping
+# are sometimes constant, e.g. christmas -> christmas which was incorrectly
+# singularized by our inflection.singularize.
+import re
+import nltk
+from nltk.corpus.reader import VERB
+import inflection
+
+_MANUAL_MATCHES = {
+    "police": "police",
+    "las": "las",
+    "vegas": "vegas",
+    "yes": "yes",
+    "jeans": "jean",
+    "hell's": "hell",
+    "domino's": "domino",
+    "morning": "morn",
+    "clothes": "cloth",
+    "are": "are",
+    "riding": "ride",
+    "leaves": "leaf",
+    "dangerous": "danger",
+    "clothing": "cloth",
+    "texting": "text",
+    "kiting": "kite",
+    "firefighters": "firefight",
+    "ties": "tie",
+    "married": "married",
+    "teething": "teeth",
+    "gloves": "glove",
+    "tennis": "tennis",
+    "dining": "dine",
+    "directions": "direct",
+    "waves": "wave",
+    "christmas": "christmas",
+    "drives": "drive",
+    "pudding": "pud",
+    "coding": "code",
+    "plating": "plate",
+    "quantas": "quanta",
+    "hornes": "horn",
+    "graves": "grave",
+    "mating": "mate",
+    "paned": "pane",
+    "alertness": "alert",
+    "sunbathing": "sunbath",
+    "tenning": "ten",
+    "wetness": "wet",
+    "urinating": "urine",
+    "sickness": "sick",
+    "braves": "brave",
+    "firefighting": "firefight",
+    "lenses": "lens",
+    "reflections": "reflect",
+    "backpackers": "backpack",
+    "eatting": "eat",
+    "designers": "design",
+    "curiousity": "curious",
+    "playfulness": "play",
+    "blindness": "blind",
+    "hawke": "hawk",
+    "tomatoe": "tomato",
+    "rodeoing": "rodeo",
+    "brightness": "bright",
+    "circuses": "circus",
+    "skateboarders": "skateboard",
+    "staring": "stare",
+    "electronics": "electron",
+    "electicity": "elect",
+    "mountainous": "mountain",
+    "socializing": "social",
+    "hamburgers": "hamburg",
+    "caves": "cave",
+    "transitions": "transit",
+    "wading": "wade",
+    "creame": "cream",
+    "toileting": "toilet",
+    "sautee": "saute",
+    "buildings": "build",
+    "belongings": "belong",
+    "stockings": "stock",
+    "walle": "wall",
+    "cumulis": "cumuli",
+    "travelers": "travel",
+    "conducter": "conduct",
+    "browsing": "brows",
+    "pooping": "poop",
+    "haircutting": "haircut",
+    "toppings": "top",
+    "hearding": "heard",
+    "sunblocker": "sunblock",
+    "bases": "base",
+    "markings": "mark",
+    "mopeds": "mope",
+    "kindergartener": "kindergarten",
+    "pies": "pie",
+    "scrapbooking": "scrapbook",
+    "couponing": "coupon",
+    "meetings": "meet",
+    "elevators": "elev",
+    "lowes": "low",
+    "men's": "men",
+    "childrens": "children",
+    "shelves": "shelve",
+    "paintings": "paint",
+    "raines": "rain",
+    "paring": "pare",
+    "expressions": "express",
+    "routes": "rout",
+    "pease": "peas",
+    "vastness": "vast",
+    "awning": "awn",
+    "boy's": "boy",
+    "drunkenness": "drunken",
+    "teasing": "teas",
+    "conferences": "confer",
+    "ripeness": "ripe",
+    "suspenders": "suspend",
+    "earnings": "earn",
+    "reporters": "report",
+    "kid's": "kid",
+    "containers": "contain",
+    "corgie": "corgi",
+    "porche": "porch",
+    "microwaves": "microwave",
+    "batter's": "batter",
+    "sadness": "sad",
+    "apartments": "apart",
+    "oxygenize": "oxygen",
+    "striping": "stripe",
+    "purring": "pure",
+    "professionals": "profession",
+    "piping": "pipe",
+    "farmer's": "farmer",
+    "potatoe": "potato",
+    "emirates": "emir",
+    "womens": "women",
+    "veteran's": "veteran",
+    "wilderness": "wilder",
+    "propellers": "propel",
+    "alpes": "alp",
+    "charioteering": "chariot",
+    "swining": "swine",
+    "illness": "ill",
+    "crepte": "crept",
+    "adhesives": "adhesive",
+    "regent's": "regent",
+    "decorations": "decor",
+    "rabbies": "rabbi",
+    "overseas": "oversea",
+    "travellers": "travel",
+    "casings": "case",
+    "smugness": "smug",
+    "doves": "dove",
+    "nationals": "nation",
+    "mustange": "mustang",
+    "ringe": "ring",
+    "gondoliere": "gondolier",
+    "vacationing": "vacate",
+    "reminders": "remind",
+    "baldness": "bald",
+    "settings": "set",
+    "glaced": "glace",
+    "coniferous": "conifer",
+    "revelations": "revel",
+    "personals": "person",
+    "daughter's": "daughter",
+    "badness": "bad",
+    "projections": "project",
+    "polarizing": "polar",
+    "vandalizers": "vandal",
+    "minerals": "miner",
+    "protesters": "protest",
+    "controllers": "control",
+    "weddings": "wed",
+    "sometimes": "sometime",
+    "earing": "ear",
+}
+
+
+class OKVQAStemmer:
+    """Stemmer to match OKVQA v1.1 procedure."""
+
+    def __init__(self):
+        self._wordnet_lemmatizer = nltk.stem.WordNetLemmatizer()
+
+    def stem(self, input_string):
+        """Apply stemming."""
+        word_and_pos = nltk.pos_tag(nltk.tokenize.word_tokenize(input_string))
+        stemmed_words = []
+        for w, p in word_and_pos:
+            if w in _MANUAL_MATCHES:
+                w = _MANUAL_MATCHES[w]
+            elif w.endswith("ing"):
+                w = self._wordnet_lemmatizer.lemmatize(w, VERB)
+            elif p.startswith("NNS") or p.startswith("NNPS"):
+                w = inflection.singularize(w)
+            stemmed_words.append(w)
+        return " ".join(stemmed_words)
+
+
+stemmer = OKVQAStemmer()
+
+
+def postprocess_ok_vqa_generation(predictions) -> str:
+    prediction = re.split("Question|Answer|Short", predictions, 1)[0]
+    prediction = re.split(", ", prediction, 1)[0]
+    prediction_stem = stemmer.stem(prediction)
+    return prediction_stem

open_flamingo/eval/rices.py

@@ -0,0 +1,95 @@
+import open_clip
+import torch
+from tqdm import tqdm
+import torch
+from utils import custom_collate_fn
+
+
+class RICES:
+    def __init__(
+        self,
+        dataset,
+        device,
+        batch_size,
+        vision_encoder_path="ViT-B-32",
+        vision_encoder_pretrained="openai",
+        cached_features=None,
+    ):
+        self.dataset = dataset
+        self.device = device
+        self.batch_size = batch_size
+
+        # Load the model and processor
+        vision_encoder, _, image_processor = open_clip.create_model_and_transforms(
+            vision_encoder_path,
+            pretrained=vision_encoder_pretrained,
+        )
+        self.model = vision_encoder.to(self.device)
+        self.image_processor = image_processor
+
+        # Precompute features
+        if cached_features is None:
+            self.features = self._precompute_features()
+        else:
+            self.features = cached_features
+
+    def _precompute_features(self):
+        features = []
+
+        # Switch to evaluation mode
+        self.model.eval()
+
+        # Set up loader
+        loader = torch.utils.data.DataLoader(
+            self.dataset,
+            batch_size=self.batch_size,
+            collate_fn=custom_collate_fn,
+        )
+
+        with torch.no_grad():
+            for batch in tqdm(
+                loader,
+                desc="Precomputing features for RICES",
+            ):
+                batch = batch["image"]
+                inputs = torch.stack(
+                    [self.image_processor(image) for image in batch]
+                ).to(self.device)
+                image_features = self.model.encode_image(inputs)
+                image_features /= image_features.norm(dim=-1, keepdim=True)
+                features.append(image_features.detach())
+
+        features = torch.cat(features)
+        return features
+
+    def find(self, batch, num_examples):
+        """
+        Get the top num_examples most similar examples to the images.
+        """
+        # Switch to evaluation mode
+        self.model.eval()
+
+        with torch.no_grad():
+            inputs = torch.stack([self.image_processor(image) for image in batch]).to(
+                self.device
+            )
+
+            # Get the feature of the input image
+            query_feature = self.model.encode_image(inputs)
+            query_feature /= query_feature.norm(dim=-1, keepdim=True)
+            query_feature = query_feature.detach().cpu()
+
+            if query_feature.ndim == 1:
+                query_feature = query_feature.unsqueeze(0)
+
+            # Compute the similarity of the input image to the precomputed features
+            similarity = (query_feature @ self.features.T).squeeze()
+
+            if similarity.ndim == 1:
+                similarity = similarity.unsqueeze(0)
+
+            # Get the indices of the 'num_examples' most similar images
+            indices = similarity.argsort(dim=-1, descending=True)[:, :num_examples]
+
+        # Return with the most similar images last
+        return [[self.dataset[i] for i in reversed(row)] for row in indices]

open_flamingo/eval/utils.py

@@ -0,0 +1,124 @@
+import numpy as np
+import torch
+import random
+import torch.nn as nn
+from contextlib import suppress
+
+
+def random_seed(seed=42, rank=0):
+    torch.manual_seed(seed + rank)
+    np.random.seed(seed + rank)
+    random.seed(seed + rank)
+
+
+def custom_collate_fn(batch):
+    """
+    Collate function for DataLoader that collates a list of dicts into a dict of lists.
+    """
+    collated_batch = {}
+    for key in batch[0].keys():
+        collated_batch[key] = [item[key] for item in batch]
+    return collated_batch
+
+
+def compute_effective_num_shots(num_shots, model_type):
+    """
+    Compute the effective number of shots for a given model type.
+    For example, following Flamingo, 0-shot OF evaluations use two text-only shots.
+    """
+    if model_type == "open_flamingo":
+        return num_shots if num_shots > 0 else 2
+    return num_shots
+
+
+def sample_batch_demos_from_query_set(query_set, num_samples, batch_size):
+    """
+    Sample random demonstrations from the query set.
+    """
+    return [random.sample(query_set, num_samples) for _ in range(batch_size)]
+
+
+def get_query_set(train_dataset, query_set_size):
+    """
+    Get a subset of the training dataset to use as the query set.
+    """
+    query_set = np.random.choice(len(train_dataset), query_set_size, replace=False)
+    return [train_dataset[i] for i in query_set]
+
+
+def prepare_eval_samples(test_dataset, num_samples, batch_size):
+    """
+    Subset the test dataset and return a DataLoader.
+    """
+    random_indices = np.random.choice(len(test_dataset), num_samples, replace=False)
+    dataset = torch.utils.data.Subset(test_dataset, random_indices)
+    sampler = torch.utils.data.distributed.DistributedSampler(dataset)
+    loader = torch.utils.data.DataLoader(
+        dataset,
+        batch_size=batch_size,
+        sampler=sampler,
+        collate_fn=custom_collate_fn,
+    )
+    return loader
+
+
+def get_indices_of_unique(x):
+    """
+    Return the indices of x that correspond to unique elements.
+    If value v is unique and two indices in x have value v, the first index is returned.
+    """
+    unique_elements = torch.unique(x)
+    first_indices = []
+    for v in unique_elements:
+        indices = torch.where(x == v)[0]
+        first_indices.append(indices[0])  # Take the first index for each unique element
+    return torch.tensor(first_indices)
+
+
+def unwrap_model(model):
+    """
+    Unwrap a model from a DataParallel or DistributedDataParallel wrapper.
+    """
+    if isinstance(model, (nn.DataParallel, nn.parallel.DistributedDataParallel)):
+        return model.module
+    else:
+        return model
+
+
+def get_predicted_classnames(logprobs, k, class_id_to_name):
+    """
+    Args:
+        - logprobs shape (B, Y) containing logprobs for each classname
+        - k: number for top-k
+        - class_id_to_name: dict mapping class index to classname
+
+    Returns:
+        - top-k predicted classnames shape (B, k) type str
+        - top-k logprobs shape (B, k) type float
+    """
+    # convert indices to classnames
+    _, predictions = torch.topk(logprobs, k=k, dim=1)  # shape (B, k)
+    predicted_classnames = [
+        [class_id_to_name[ix] for ix in item] for item in predictions.tolist()
+    ]
+    predicted_logprobs = torch.gather(logprobs, 1, predictions)
+    return predicted_classnames, predicted_logprobs
+
+
+def get_cast_dtype(precision: str):
+    cast_dtype = None
+    if precision == "bf16":
+        cast_dtype = torch.bfloat16
+    elif precision == "fp16":
+        cast_dtype = torch.float16
+    return cast_dtype
+
+
+def get_autocast(precision):
+    if precision == "amp":
+        return torch.cuda.amp.autocast
+    elif precision == "amp_bfloat16" or precision == "amp_bf16":
+        # amp_bfloat16 is more stable than amp float16 for clip training
+        return lambda: torch.cuda.amp.autocast(dtype=torch.bfloat16)
+    else:
+        return suppress

open_flamingo/eval/vqa_metric.py

@@ -0,0 +1,560 @@
+import copy
+import datetime
+import json
+import os
+import random
+import re
+import sys
+
+# Interface for accessing the VQA dataset.
+
+# This code is based on the code written by Tsung-Yi Lin for MSCOCO Python API available at the following link:
+# (https://github.com/pdollar/coco/blob/master/PythonAPI/pycocotools/coco.py).
+
+# The following functions are defined:
+#  VQA        - VQA class that loads VQA annotation file and prepares data structures.
+#  getQuesIds - Get question ids that satisfy given filter conditions.
+#  getImgIds  - Get image ids that satisfy given filter conditions.
+#  loadQA     - Load questions and answers with the specified question ids.
+#  showQA     - Display the specified questions and answers.
+#  loadRes    - Load result file and create result object.
+
+# Help on each function can be accessed by: "help(COCO.function)"
+
+
+class VQA:
+    def __init__(self, annotation_file=None, question_file=None):
+        """
+        Constructor of VQA helper class for reading and visualizing questions and answers.
+        :param annotation_file (str): location of VQA annotation file
+        :return:
+        """
+        # load dataset
+        self.dataset = {}
+        self.questions = {}
+        self.qa = {}
+        self.qqa = {}
+        self.imgToQA = {}
+        if not annotation_file == None and not question_file == None:
+            print("loading VQA annotations and questions into memory...")
+            time_t = datetime.datetime.utcnow()
+            dataset = json.load(open(annotation_file, "r"))
+            questions = json.load(open(question_file, "r"))
+            print(datetime.datetime.utcnow() - time_t)
+            self.dataset = dataset
+            self.questions = questions
+            self.createIndex()
+
+    def createIndex(self):
+        # create index
+        print("creating index...")
+        imgToQA = {ann["image_id"]: [] for ann in self.dataset["annotations"]}
+        qa = {ann["question_id"]: [] for ann in self.dataset["annotations"]}
+        qqa = {ann["question_id"]: [] for ann in self.dataset["annotations"]}
+        for ann in self.dataset["annotations"]:
+            imgToQA[ann["image_id"]] += [ann]
+            qa[ann["question_id"]] = ann
+        for ques in self.questions["questions"]:
+            qqa[ques["question_id"]] = ques
+        print("index created!")
+
+        # create class members
+        self.qa = qa
+        self.qqa = qqa
+        self.imgToQA = imgToQA
+
+    def info(self):
+        """
+        Print information about the VQA annotation file.
+        :return:
+        """
+        for key, value in self.dataset["info"].items():
+            print("%s: %s" % (key, value))
+
+    def getQuesIds(self, imgIds=[], quesTypes=[], ansTypes=[]):
+        """
+        Get question ids that satisfy given filter conditions. default skips that filter
+        :param 	imgIds    (int array)   : get question ids for given imgs
+                        quesTypes (str array)   : get question ids for given question types
+                        ansTypes  (str array)   : get question ids for given answer types
+        :return:    ids   (int array)   : integer array of question ids
+        """
+        imgIds = imgIds if type(imgIds) == list else [imgIds]
+        quesTypes = quesTypes if type(quesTypes) == list else [quesTypes]
+        ansTypes = ansTypes if type(ansTypes) == list else [ansTypes]
+
+        if len(imgIds) == len(quesTypes) == len(ansTypes) == 0:
+            anns = self.dataset["annotations"]
+        else:
+            if not len(imgIds) == 0:
+                anns = sum(
+                    [self.imgToQA[imgId] for imgId in imgIds if imgId in self.imgToQA],
+                    [],
+                )
+            else:
+                anns = self.dataset["annotations"]
+            anns = (
+                anns
+                if len(quesTypes) == 0
+                else [ann for ann in anns if ann["question_type"] in quesTypes]
+            )
+            anns = (
+                anns
+                if len(ansTypes) == 0
+                else [ann for ann in anns if ann["answer_type"] in ansTypes]
+            )
+        ids = [ann["question_id"] for ann in anns]
+        return ids
+
+    def getImgIds(self, quesIds=[], quesTypes=[], ansTypes=[]):
+        """
+         Get image ids that satisfy given filter conditions. default skips that filter
+         :param quesIds   (int array)   : get image ids for given question ids
+        quesTypes (str array)   : get image ids for given question types
+        ansTypes  (str array)   : get image ids for given answer types
+         :return: ids     (int array)   : integer array of image ids
+        """
+        quesIds = quesIds if type(quesIds) == list else [quesIds]
+        quesTypes = quesTypes if type(quesTypes) == list else [quesTypes]
+        ansTypes = ansTypes if type(ansTypes) == list else [ansTypes]
+
+        if len(quesIds) == len(quesTypes) == len(ansTypes) == 0:
+            anns = self.dataset["annotations"]
+        else:
+            if not len(quesIds) == 0:
+                anns = sum(
+                    [self.qa[quesId] for quesId in quesIds if quesId in self.qa], []
+                )
+            else:
+                anns = self.dataset["annotations"]
+            anns = (
+                anns
+                if len(quesTypes) == 0
+                else [ann for ann in anns if ann["question_type"] in quesTypes]
+            )
+            anns = (
+                anns
+                if len(ansTypes) == 0
+                else [ann for ann in anns if ann["answer_type"] in ansTypes]
+            )
+        ids = [ann["image_id"] for ann in anns]
+        return ids
+
+    def loadQA(self, ids=[]):
+        """
+        Load questions and answers with the specified question ids.
+        :param ids (int array)       : integer ids specifying question ids
+        :return: qa (object array)   : loaded qa objects
+        """
+        if type(ids) == list:
+            return [self.qa[id] for id in ids]
+        elif type(ids) == int:
+            return [self.qa[ids]]
+
+    def showQA(self, anns):
+        """
+        Display the specified annotations.
+        :param anns (array of object): annotations to display
+        :return: None
+        """
+        if len(anns) == 0:
+            return 0
+        for ann in anns:
+            quesId = ann["question_id"]
+            print("Question: %s" % (self.qqa[quesId]["question"]))
+            for ans in ann["answers"]:
+                print("Answer %d: %s" % (ans["answer_id"], ans["answer"]))
+
+    def loadRes(self, resFile, quesFile):
+        """
+        Load result file and return a result object.
+        :param   resFile (str)     : file name of result file
+        :return: res (obj)         : result api object
+        """
+        res = VQA()
+        res.questions = json.load(open(quesFile))
+        res.dataset["info"] = copy.deepcopy(self.questions["info"])
+        res.dataset["task_type"] = copy.deepcopy(self.questions["task_type"])
+        res.dataset["data_type"] = copy.deepcopy(self.questions["data_type"])
+        res.dataset["data_subtype"] = copy.deepcopy(self.questions["data_subtype"])
+        res.dataset["license"] = copy.deepcopy(self.questions["license"])
+
+        print("Loading and preparing results...     ")
+        time_t = datetime.datetime.utcnow()
+        anns = json.load(open(resFile))
+        assert type(anns) == list, "results is not an array of objects"
+        annsQuesIds = [ann["question_id"] for ann in anns]
+        # print set of question ids that do not have corresponding annotations
+
+        # assert set(annsQuesIds) == set(self.getQuesIds()), \
+        # 'Results do not correspond to current VQA set. Either the results do not have predictions for all question ids in annotation file or there is atleast one question id that does not belong to the question ids in the annotation file.'
+        for ann in anns:
+            quesId = ann["question_id"]
+            if res.dataset["task_type"] == "Multiple Choice":
+                assert (
+                    ann["answer"] in self.qqa[quesId]["multiple_choices"]
+                ), "predicted answer is not one of the multiple choices"
+            qaAnn = self.qa[quesId]
+            ann["image_id"] = qaAnn["image_id"]
+            ann["question_type"] = qaAnn["question_type"]
+            if "answer_type" in ann:
+                ann["answer_type"] = qaAnn["answer_type"]
+        print(
+            "DONE (t=%0.2fs)" % ((datetime.datetime.utcnow() - time_t).total_seconds())
+        )
+
+        res.dataset["annotations"] = anns
+        res.createIndex()
+        return res
+
+
+class VQAEval:
+    def __init__(self, vqa, vqaRes, n=2):
+        self.n = n
+        self.accuracy = {}
+        self.evalQA = {}
+        self.evalQuesType = {}
+        self.evalAnsType = {}
+        self.vqa = vqa
+        self.vqaRes = vqaRes
+        if not vqa is None and not vqaRes is None:
+            self.params = {"question_id": vqaRes.getQuesIds()}
+        self.contractions = {
+            "aint": "ain't",
+            "arent": "aren't",
+            "cant": "can't",
+            "couldve": "could've",
+            "couldnt": "couldn't",
+            "couldn'tve": "couldn't've",
+            "couldnt've": "couldn't've",
+            "didnt": "didn't",
+            "doesnt": "doesn't",
+            "dont": "don't",
+            "hadnt": "hadn't",
+            "hadnt've": "hadn't've",
+            "hadn'tve": "hadn't've",
+            "hasnt": "hasn't",
+            "havent": "haven't",
+            "hed": "he'd",
+            "hed've": "he'd've",
+            "he'dve": "he'd've",
+            "hes": "he's",
+            "howd": "how'd",
+            "howll": "how'll",
+            "hows": "how's",
+            "Id've": "I'd've",
+            "I'dve": "I'd've",
+            "Im": "I'm",
+            "Ive": "I've",
+            "isnt": "isn't",
+            "itd": "it'd",
+            "itd've": "it'd've",
+            "it'dve": "it'd've",
+            "itll": "it'll",
+            "let's": "let's",
+            "maam": "ma'am",
+            "mightnt": "mightn't",
+            "mightnt've": "mightn't've",
+            "mightn'tve": "mightn't've",
+            "mightve": "might've",
+            "mustnt": "mustn't",
+            "mustve": "must've",
+            "neednt": "needn't",
+            "notve": "not've",
+            "oclock": "o'clock",
+            "oughtnt": "oughtn't",
+            "ow's'at": "'ow's'at",
+            "'ows'at": "'ow's'at",
+            "'ow'sat": "'ow's'at",
+            "shant": "shan't",
+            "shed've": "she'd've",
+            "she'dve": "she'd've",
+            "she's": "she's",
+            "shouldve": "should've",
+            "shouldnt": "shouldn't",
+            "shouldnt've": "shouldn't've",
+            "shouldn'tve": "shouldn't've",
+            "somebody'd": "somebodyd",
+            "somebodyd've": "somebody'd've",
+            "somebody'dve": "somebody'd've",
+            "somebodyll": "somebody'll",
+            "somebodys": "somebody's",
+            "someoned": "someone'd",
+            "someoned've": "someone'd've",
+            "someone'dve": "someone'd've",
+            "someonell": "someone'll",
+            "someones": "someone's",
+            "somethingd": "something'd",
+            "somethingd've": "something'd've",
+            "something'dve": "something'd've",
+            "somethingll": "something'll",
+            "thats": "that's",
+            "thered": "there'd",
+            "thered've": "there'd've",
+            "there'dve": "there'd've",
+            "therere": "there're",
+            "theres": "there's",
+            "theyd": "they'd",
+            "theyd've": "they'd've",
+            "they'dve": "they'd've",
+            "theyll": "they'll",
+            "theyre": "they're",
+            "theyve": "they've",
+            "twas": "'twas",
+            "wasnt": "wasn't",
+            "wed've": "we'd've",
+            "we'dve": "we'd've",
+            "weve": "we've",
+            "werent": "weren't",
+            "whatll": "what'll",
+            "whatre": "what're",
+            "whats": "what's",
+            "whatve": "what've",
+            "whens": "when's",
+            "whered": "where'd",
+            "wheres": "where's",
+            "whereve": "where've",
+            "whod": "who'd",
+            "whod've": "who'd've",
+            "who'dve": "who'd've",
+            "wholl": "who'll",
+            "whos": "who's",
+            "whove": "who've",
+            "whyll": "why'll",
+            "whyre": "why're",
+            "whys": "why's",
+            "wont": "won't",
+            "wouldve": "would've",
+            "wouldnt": "wouldn't",
+            "wouldnt've": "wouldn't've",
+            "wouldn'tve": "wouldn't've",
+            "yall": "y'all",
+            "yall'll": "y'all'll",
+            "y'allll": "y'all'll",
+            "yall'd've": "y'all'd've",
+            "y'alld've": "y'all'd've",
+            "y'all'dve": "y'all'd've",
+            "youd": "you'd",
+            "youd've": "you'd've",
+            "you'dve": "you'd've",
+            "youll": "you'll",
+            "youre": "you're",
+            "youve": "you've",
+        }
+        self.manualMap = {
+            "none": "0",
+            "zero": "0",
+            "one": "1",
+            "two": "2",
+            "three": "3",
+            "four": "4",
+            "five": "5",
+            "six": "6",
+            "seven": "7",
+            "eight": "8",
+            "nine": "9",
+            "ten": "10",
+        }
+        self.articles = ["a", "an", "the"]
+
+        self.periodStrip = re.compile("(?!<=\d)(\.)(?!\d)")
+        self.commaStrip = re.compile("(\d)(\,)(\d)")
+        self.punct = [
+            ";",
+            r"/",
+            "[",
+            "]",
+            '"',
+            "{",
+            "}",
+            "(",
+            ")",
+            "=",
+            "+",
+            "\\",
+            "_",
+            "-",
+            ">",
+            "<",
+            "@",
+            "`",
+            ",",
+            "?",
+            "!",
+        ]
+
+    def evaluate(self, quesIds=None):
+        if quesIds == None:
+            quesIds = [quesId for quesId in self.params["question_id"]]
+        gts = {}
+        res = {}
+        for quesId in quesIds:
+            gts[quesId] = self.vqa.qa[quesId]
+            res[quesId] = self.vqaRes.qa[quesId]
+
+        # =================================================
+        # Compute accuracy
+        # =================================================
+        accQA = []
+        accQuesType = {}
+        accAnsType = {}
+        print("computing accuracy")
+        step = 0
+        for quesId in quesIds:
+            for ansDic in gts[quesId]["answers"]:
+                ansDic["answer"] = ansDic["answer"].replace("\n", " ")
+                ansDic["answer"] = ansDic["answer"].replace("\t", " ")
+                ansDic["answer"] = ansDic["answer"].strip()
+            resAns = res[quesId]["answer"]
+            resAns = resAns.replace("\n", " ")
+            resAns = resAns.replace("\t", " ")
+            resAns = resAns.strip()
+            resAns = self.processPunctuation(resAns)
+            resAns = self.processDigitArticle(resAns)
+            gtAcc = []
+
+            for ansDic in gts[quesId]["answers"]:
+                ansDic["answer"] = self.processPunctuation(ansDic["answer"])
+                ansDic["answer"] = self.processDigitArticle(ansDic["answer"])
+
+            for gtAnsDatum in gts[quesId]["answers"]:
+                otherGTAns = [
+                    item for item in gts[quesId]["answers"] if item != gtAnsDatum
+                ]
+                matchingAns = [item for item in otherGTAns if item["answer"] == resAns]
+                acc = min(1, float(len(matchingAns)) / 3)
+                gtAcc.append(acc)
+            quesType = gts[quesId]["question_type"]
+            ansType = (
+                gts[quesId]["answer_type"] if "answer_type" in gts[quesId] else "other"
+            )
+            avgGTAcc = float(sum(gtAcc)) / len(gtAcc)
+            accQA.append(avgGTAcc)
+            if quesType not in accQuesType:
+                accQuesType[quesType] = []
+            accQuesType[quesType].append(avgGTAcc)
+            if ansType not in accAnsType:
+                accAnsType[ansType] = []
+            accAnsType[ansType].append(avgGTAcc)
+            self.setEvalQA(quesId, avgGTAcc)
+            self.setEvalQuesType(quesId, quesType, avgGTAcc)
+            self.setEvalAnsType(quesId, ansType, avgGTAcc)
+            if step % 100 == 0:
+                self.updateProgress(step / float(len(quesIds)))
+            step = step + 1
+
+        self.setAccuracy(accQA, accQuesType, accAnsType)
+        print("Done computing accuracy")
+
+    def processPunctuation(self, inText):
+        outText = inText
+        for p in self.punct:
+            if (p + " " in inText or " " + p in inText) or (
+                re.search(self.commaStrip, inText) != None
+            ):
+                outText = outText.replace(p, "")
+            else:
+                outText = outText.replace(p, " ")
+        outText = self.periodStrip.sub("", outText, re.UNICODE)
+        return outText
+
+    def processDigitArticle(self, inText):
+        outText = []
+        tempText = inText.lower().split()
+        for word in tempText:
+            word = self.manualMap.setdefault(word, word)
+            if word not in self.articles:
+                outText.append(word)
+            else:
+                pass
+        for wordId, word in enumerate(outText):
+            if word in self.contractions:
+                outText[wordId] = self.contractions[word]
+        outText = " ".join(outText)
+        return outText
+
+    def setAccuracy(self, accQA, accQuesType, accAnsType):
+        self.accuracy["overall"] = round(100 * float(sum(accQA)) / len(accQA), self.n)
+        self.accuracy["perQuestionType"] = {
+            quesType: round(
+                100 * float(sum(accQuesType[quesType])) / len(accQuesType[quesType]),
+                self.n,
+            )
+            for quesType in accQuesType
+        }
+        self.accuracy["perAnswerType"] = {
+            ansType: round(
+                100 * float(sum(accAnsType[ansType])) / len(accAnsType[ansType]), self.n
+            )
+            for ansType in accAnsType
+        }
+
+    def setEvalQA(self, quesId, acc):
+        self.evalQA[quesId] = round(100 * acc, self.n)
+
+    def setEvalQuesType(self, quesId, quesType, acc):
+        if quesType not in self.evalQuesType:
+            self.evalQuesType[quesType] = {}
+        self.evalQuesType[quesType][quesId] = round(100 * acc, self.n)
+
+    def setEvalAnsType(self, quesId, ansType, acc):
+        if ansType not in self.evalAnsType:
+            self.evalAnsType[ansType] = {}
+        self.evalAnsType[ansType][quesId] = round(100 * acc, self.n)
+
+    def updateProgress(self, progress):
+        barLength = 20
+        status = ""
+        if isinstance(progress, int):
+            progress = float(progress)
+        if not isinstance(progress, float):
+            progress = 0
+            status = "error: progress var must be float\r\n"
+        if progress < 0:
+            progress = 0
+            status = "Halt...\r\n"
+        if progress >= 1:
+            progress = 1
+            status = "Done...\r\n"
+        block = int(round(barLength * progress))
+        text = "\rFinshed Percent: [{0}] {1}% {2}".format(
+            "#" * block + "-" * (barLength - block), int(progress * 100), status
+        )
+        sys.stdout.write(text)
+        sys.stdout.flush()
+
+
+def compute_vqa_accuracy(result_json_path, question_json_path, annotation_json_path):
+    """Compute the VQA accuracy metric.
+
+    Args:
+        result_json_path (str): Path to the json file with model outputs
+        question_json_path (str): Path to the json file with questions
+        annotation_json_path (str): Path to the json file with annotations
+
+    Returns:
+        float: VQA accuracy
+    """
+
+    # create vqa object and vqaRes object
+    vqa = VQA(annotation_json_path, question_json_path)
+    vqaRes = vqa.loadRes(result_json_path, question_json_path)
+
+    # create vqaEval object by taking vqa and vqaRes
+    # n is precision of accuracy (number of places after decimal), default is 2
+    vqaEval = VQAEval(vqa, vqaRes, n=2)
+
+    # evaluate results
+    """
+    If you have a list of question ids on which you would like to evaluate your results, pass it as a list to below function
+    By default it uses all the question ids in annotation file
+    """
+    vqaEval.evaluate()
+
+    return vqaEval.accuracy["overall"]
+
+
+def postprocess_vqa_generation(predictions):
+    answer = re.split("Question|Answer|Short", predictions, 1)[0]
+    answer = re.split(", ", answer, 1)[0]
+    return answer

open_flamingo/scripts/cache_rices_features.py

@@ -0,0 +1,370 @@
+"""
+Cache CLIP features for all images in training split in preparation for RICES
+"""
+import argparse
+import sys
+import os
+
+sys.path.append(
+    os.path.join(
+        os.path.dirname(os.path.abspath(__file__)),
+        "..",
+    )
+)
+from eval.rices import RICES
+from eval.eval_datasets import (
+    CaptionDataset,
+    VQADataset,
+    ImageNetDataset,
+    HatefulMemesDataset,
+)
+import os
+import torch
+
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    "--output_dir",
+    type=str,
+    required=True,
+    help="Directory to save the cached features.",
+)
+parser.add_argument("--vision_encoder_path", default="ViT-L-14", type=str)
+parser.add_argument("--vision_encoder_pretrained", default="openai", type=str)
+parser.add_argument("--batch_size", default=256)
+
+# Per-dataset flags
+parser.add_argument(
+    "--eval_coco",
+    action="store_true",
+    default=False,
+    help="Whether to cache COCO.",
+)
+parser.add_argument(
+    "--eval_vqav2",
+    action="store_true",
+    default=False,
+    help="Whether to cache VQAV2.",
+)
+parser.add_argument(
+    "--eval_ok_vqa",
+    action="store_true",
+    default=False,
+    help="Whether to cache OK-VQA.",
+)
+parser.add_argument(
+    "--eval_vizwiz",
+    action="store_true",
+    default=False,
+    help="Whether to cache VizWiz.",
+)
+parser.add_argument(
+    "--eval_textvqa",
+    action="store_true",
+    default=False,
+    help="Whether to cache TextVQA.",
+)
+parser.add_argument(
+    "--eval_imagenet",
+    action="store_true",
+    default=False,
+    help="Whether to cache ImageNet.",
+)
+parser.add_argument(
+    "--eval_flickr30",
+    action="store_true",
+    default=False,
+    help="Whether to cache Flickr30.",
+)
+parser.add_argument(
+    "--eval_hateful_memes",
+    action="store_true",
+    default=False,
+    help="Whether to cache Hateful Memes.",
+)
+
+# Dataset arguments
+
+## Flickr30 Dataset
+parser.add_argument(
+    "--flickr_image_dir_path",
+    type=str,
+    help="Path to the flickr30/flickr30k_images directory.",
+    default=None,
+)
+parser.add_argument(
+    "--flickr_karpathy_json_path",
+    type=str,
+    help="Path to the dataset_flickr30k.json file.",
+    default=None,
+)
+parser.add_argument(
+    "--flickr_annotations_json_path",
+    type=str,
+    help="Path to the dataset_flickr30k_coco_style.json file.",
+)
+## COCO Dataset
+parser.add_argument(
+    "--coco_train_image_dir_path",
+    type=str,
+    default=None,
+)
+parser.add_argument(
+    "--coco_val_image_dir_path",
+    type=str,
+    default=None,
+)
+parser.add_argument(
+    "--coco_karpathy_json_path",
+    type=str,
+    default=None,
+)
+parser.add_argument(
+    "--coco_annotations_json_path",
+    type=str,
+    default=None,
+)
+
+## VQAV2 Dataset
+parser.add_argument(
+    "--vqav2_train_image_dir_path",
+    type=str,
+    default=None,
+)
+parser.add_argument(
+    "--vqav2_train_questions_json_path",
+    type=str,
+    default=None,
+)
+parser.add_argument(
+    "--vqav2_train_annotations_json_path",
+    type=str,
+    default=None,
+)
+
+## OK-VQA Dataset
+parser.add_argument(
+    "--ok_vqa_train_image_dir_path",
+    type=str,
+    help="Path to the vqav2/train2014 directory.",
+    default=None,
+)
+parser.add_argument(
+    "--ok_vqa_train_questions_json_path",
+    type=str,
+    help="Path to the v2_OpenEnded_mscoco_train2014_questions.json file.",
+    default=None,
+)
+parser.add_argument(
+    "--ok_vqa_train_annotations_json_path",
+    type=str,
+    help="Path to the v2_mscoco_train2014_annotations.json file.",
+    default=None,
+)
+
+## VizWiz Dataset
+parser.add_argument(
+    "--vizwiz_train_image_dir_path",
+    type=str,
+    help="Path to the vizwiz train images directory.",
+    default=None,
+)
+parser.add_argument(
+    "--vizwiz_train_questions_json_path",
+    type=str,
+    help="Path to the vizwiz questions json file.",
+    default=None,
+)
+parser.add_argument(
+    "--vizwiz_train_annotations_json_path",
+    type=str,
+    help="Path to the vizwiz annotations json file.",
+    default=None,
+)
+
+# TextVQA Dataset
+parser.add_argument(
+    "--textvqa_image_dir_path",
+    type=str,
+    help="Path to the textvqa images directory.",
+    default=None,
+)
+parser.add_argument(
+    "--textvqa_train_questions_json_path",
+    type=str,
+    help="Path to the textvqa questions json file.",
+    default=None,
+)
+parser.add_argument(
+    "--textvqa_train_annotations_json_path",
+    type=str,
+    help="Path to the textvqa annotations json file.",
+    default=None,
+)
+
+
+## Imagenet dataset
+parser.add_argument("--imagenet_root", type=str, default="/tmp")
+
+## Hateful Memes dataset
+parser.add_argument(
+    "--hateful_memes_image_dir_path",
+    type=str,
+    default=None,
+)
+parser.add_argument(
+    "--hateful_memes_train_annotations_json_path",
+    type=str,
+    default=None,
+)
+
+
+def main():
+    args, leftovers = parser.parse_known_args()
+    device_id = torch.cuda.current_device() if torch.cuda.is_available() else "cpu"
+    if args.eval_flickr30:
+        print("Caching Flickr30k...")
+        train_dataset = CaptionDataset(
+            image_train_dir_path=args.flickr_image_dir_path,
+            image_val_dir_path=None,
+            annotations_path=args.flickr_karpathy_json_path,
+            is_train=True,
+            dataset_name="flickr",
+        )
+        rices_dataset = RICES(
+            train_dataset,
+            device_id,
+            args.batch_size,
+            vision_encoder_path=args.vision_encoder_path,
+            vision_encoder_pretrained=args.vision_encoder_pretrained,
+        )
+        torch.save(
+            rices_dataset.features,
+            os.path.join(args.output_dir, "flickr30.pkl"),
+        )
+
+    if args.eval_coco:
+        print("Caching COCO...")
+        train_dataset = CaptionDataset(
+            image_train_dir_path=args.coco_train_image_dir_path,
+            image_val_dir_path=args.coco_val_image_dir_path,
+            annotations_path=args.coco_karpathy_json_path,
+            is_train=True,
+            dataset_name="coco",
+        )
+        rices_dataset = RICES(
+            train_dataset,
+            device_id,
+            args.batch_size,
+            vision_encoder_path=args.vision_encoder_path,
+            vision_encoder_pretrained=args.vision_encoder_pretrained,
+        )
+        torch.save(
+            rices_dataset.features,
+            os.path.join(args.output_dir, "coco.pkl"),
+        )
+
+    if args.eval_ok_vqa:
+        print("Caching OK-VQA...")
+        train_dataset = VQADataset(
+            image_dir_path=args.ok_vqa_train_image_dir_path,
+            question_path=args.ok_vqa_train_questions_json_path,
+            annotations_path=args.ok_vqa_train_annotations_json_path,
+            is_train=True,
+            dataset_name="ok_vqa",
+        )
+        rices_dataset = RICES(
+            train_dataset,
+            device_id,
+            args.batch_size,
+            vision_encoder_path=args.vision_encoder_path,
+            vision_encoder_pretrained=args.vision_encoder_pretrained,
+        )
+        torch.save(
+            rices_dataset.features,
+            os.path.join(args.output_dir, "ok_vqa.pkl"),
+        )
+
+    if args.eval_vizwiz:
+        print("Caching VizWiz...")
+        train_dataset = VQADataset(
+            image_dir_path=args.vizwiz_train_image_dir_path,
+            question_path=args.vizwiz_train_questions_json_path,
+            annotations_path=args.vizwiz_train_annotations_json_path,
+            is_train=True,
+            dataset_name="vizwiz",
+        )
+        rices_dataset = RICES(
+            train_dataset,
+            device_id,
+            args.batch_size,
+            vision_encoder_path=args.vision_encoder_path,
+            vision_encoder_pretrained=args.vision_encoder_pretrained,
+        )
+        torch.save(
+            rices_dataset.features,
+            os.path.join(args.output_dir, "vizwiz.pkl"),
+        )
+
+    if args.eval_vqav2:
+        print("Caching VQAv2...")
+        train_dataset = VQADataset(
+            image_dir_path=args.vqav2_train_image_dir_path,
+            question_path=args.vqav2_train_questions_json_path,
+            annotations_path=args.vqav2_train_annotations_json_path,
+            is_train=True,
+            dataset_name="vqav2",
+        )
+        rices_dataset = RICES(
+            train_dataset,
+            device_id,
+            args.batch_size,
+            vision_encoder_path=args.vision_encoder_path,
+            vision_encoder_pretrained=args.vision_encoder_pretrained,
+        )
+        torch.save(
+            rices_dataset.features,
+            os.path.join(args.output_dir, "vqav2.pkl"),
+        )
+
+    if args.eval_textvqa:
+        print("Caching TextVQA...")
+        train_dataset = VQADataset(
+            image_dir_path=args.textvqa_image_dir_path,
+            question_path=args.textvqa_train_questions_json_path,
+            annotations_path=args.textvqa_train_annotations_json_path,
+            is_train=True,
+            dataset_name="textvqa",
+        )
+        rices_dataset = RICES(
+            train_dataset,
+            device_id,
+            args.batch_size,
+            vision_encoder_path=args.vision_encoder_path,
+            vision_encoder_pretrained=args.vision_encoder_pretrained,
+        )
+        torch.save(
+            rices_dataset.features,
+            os.path.join(args.output_dir, "textvqa.pkl"),
+        )
+
+    if args.eval_hateful_memes:
+        print("Caching Hateful Memes...")
+        train_dataset = HatefulMemesDataset(
+            image_dir_path=args.hateful_memes_image_dir_path,
+            annotations_path=args.hateful_memes_train_annotations_json_path,
+        )
+        rices_dataset = RICES(
+            train_dataset,
+            device_id,
+            args.batch_size,
+            vision_encoder_path=args.vision_encoder_path,
+            vision_encoder_pretrained=args.vision_encoder_pretrained,
+        )
+        torch.save(
+            rices_dataset.features,
+            os.path.join(args.output_dir, "hateful_memes.pkl"),
+        )
+
+
+if __name__ == "__main__":
+    main()

open_flamingo/scripts/convert_mmc4_to_wds.py

@@ -0,0 +1,85 @@
+import argparse
+import json
+import os
+import uuid
+import zipfile
+from PIL import Image
+import base64
+from io import BytesIO
+
+import braceexpand
+import webdataset as wds
+
+arg_parser = argparse.ArgumentParser()
+arg_parser.add_argument(
+    "--output_dir",
+    type=str,
+    help="Pass in the directory where the output shards (as tar files) will be written to.",
+)
+arg_parser.add_argument(
+    "--zip_files",
+    type=str,
+    help="Pass in a list of MMC4 shards in the format path_to_shard/shard_{0..23098}.zip",
+)
+arg_parser.add_argument(
+    "--image_dir",
+    type=str,
+    help="Pass in the directory where the images have been downloaded to.",
+)
+arg_parser.add_argument(
+    "--num_files_per_shard",
+    type=int,
+    default=1000,
+)
+args = arg_parser.parse_args()
+
+
+def main():
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    doc_shards = list(braceexpand.braceexpand(args.zip_files))
+
+    with wds.ShardWriter(args.output_dir + "/%09d.tar") as sink:
+        for idx in range(len(doc_shards)):
+            # Open the ZIP archive and extract the JSON file
+            with zipfile.ZipFile(doc_shards[idx], "r") as zip_file:
+                # Assumes the JSON file is the first file in the archive
+                json_filename = zip_file.namelist()[0]
+                with zip_file.open(json_filename, "r") as json_file:
+                    for sample_data in json_file:
+                        # get image names from json
+                        sample_data = json.loads(sample_data)
+                        image_info = sample_data["image_info"]
+                        image_names = [image["image_name"] for image in image_info]
+
+                        # Add each image to the tar file
+                        for img_idx, image_name in enumerate(image_names):
+                            try:
+                                # load image
+                                img = Image.open(
+                                    os.path.join(args.image_dir, str(idx), image_name)
+                                ).convert("RGB")
+                                buffered = BytesIO()
+                                img.save(buffered, format="JPEG")
+                                img_str = base64.b64encode(buffered.getvalue())
+
+                                # convert to base64
+                                sample_data["image_info"][img_idx][
+                                    "image_base64"
+                                ] = img_str.decode("utf-8")
+                            except FileNotFoundError:
+                                print(
+                                    f"Did not find {image_name} downloaded. This can happen if the url is now 404."
+                                )
+                            except Exception as e:
+                                print(f"Error processing {image_name}: {e}")
+
+                        key_str = uuid.uuid4().hex
+                        sink.write({"__key__": key_str, "json": sample_data})
+
+            if (idx + 1) % args.num_files_per_shard == 0:
+                sink.next_stream()
+
+
+if __name__ == "__main__":
+    main()

open_flamingo/scripts/fill_vqa_testdev_results.py

@@ -0,0 +1,142 @@
+"""
+Helper scripts to prepare a vqa test-dev evaluation for EvalAI submission.
+Note: EvalAI requires VQAv2 submissions to have predictions for all the questions in the test2015 set, not just the test-dev set.
+Given a json with a subset of the vqa questions, fill in the rest of the questions with an empty string as the model prediction.
+"""
+import json
+import sys
+import os
+
+sys.path.append(
+    os.path.join(
+        os.path.dirname(os.path.abspath(__file__)),
+        "..",
+    )
+)
+from eval.vqa_metric import VQAEval
+
+postprocessor = VQAEval(None, None)
+
+
+def fill_vizwiz_test_json(
+    input_path,
+    output_path,
+    vqa_test_questions_json_path,
+):
+    # read the input json and build a set with all question_ids
+    with open(input_path, "r") as f:
+        input_json = json.load(f)
+
+    # postprocess answers
+    question_id_to_answer = {}
+    for q in input_json:
+        resAns = q["answer"]
+        resAns = resAns.replace("\n", " ")
+        resAns = resAns.replace("\t", " ")
+        resAns = resAns.strip()
+        resAns = postprocessor.processPunctuation(resAns)
+        resAns = postprocessor.processDigitArticle(resAns)
+        question_id_to_answer[q["question_id"]] = resAns
+
+    # read the vqa test json to get all the qustion_ids that need to be filled
+    with open(vqa_test_questions_json_path, "r") as f:
+        vqa_test_json = json.load(f)
+    vqa_test_json = vqa_test_json["questions"]
+
+    # if the question_id is not in the set, add it to the copy of the input json with an empty string as the answer
+    output_json = []
+    for q in vqa_test_json:
+        output_json.append(
+            {
+                "image": q["image_id"],
+                "answer": question_id_to_answer.get(q["question_id"], ""),
+            }
+        )
+
+    # write the json to the output path
+    with open(output_path, "w") as f:
+        json.dump(output_json, f)
+
+
+def fill_vqav2_test_json(
+    input_path,
+    output_path,
+    vqa_test_questions_json_path,
+):
+    # read the input json and build a set with all question_ids
+    with open(input_path, "r") as f:
+        input_json = json.load(f)
+    question_ids = set()
+    for q in input_json:
+        question_ids.add(q["question_id"])
+
+    # make a copy of the input json
+    output_json = []
+    for q in input_json:
+        resAns = q["answer"]
+        resAns = resAns.replace("\n", " ")
+        resAns = resAns.replace("\t", " ")
+        resAns = resAns.strip()
+        resAns = postprocessor.processPunctuation(resAns)
+        resAns = postprocessor.processDigitArticle(resAns)
+        q["answer"] = resAns
+        output_json.append(q)
+
+    # read the vqa test json to get all the qustion_ids that need to be filled
+    with open(vqa_test_questions_json_path, "r") as f:
+        vqa_test_json = json.load(f)
+    vqa_test_json = vqa_test_json["questions"]
+
+    # if the question_id is not in the set, add it to the copy of the input json with an empty string as the answer
+    for q in vqa_test_json:
+        if q["question_id"] not in question_ids:
+            output_json.append(
+                {
+                    "question_id": q["question_id"],
+                    "answer": "",
+                }
+            )
+
+    # write the json to the output path
+    with open(output_path, "w") as f:
+        json.dump(output_json, f)
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--dataset",
+        type=str,
+        choices=["vqav2", "vizwiz"],
+    )
+    parser.add_argument(
+        "--input_path",
+        type=str,
+        help="Path to the json file with the subset of the vqa test-dev questions.",
+    )
+    parser.add_argument(
+        "--vqa_test_questions_json_path",
+        type=str,
+        help="Path to the json file with all the vqa test questions.",
+    )
+    parser.add_argument(
+        "--output_path",
+        type=str,
+        help="Path to store the filled json.",
+    )
+    args = parser.parse_args()
+
+    if args.dataset == "vqav2":
+        fill_vqav2_test_json(
+            args.input_path,
+            args.output_path,
+            args.vqa_test_questions_json_path,
+        )
+    else:
+        fill_vizwiz_test_json(
+            args.input_path,
+            args.output_path,
+            args.vqa_test_questions_json_path,
+        )

open_flamingo/train/README.md

@@ -0,0 +1,65 @@
+# OpenFlamingo Training
+To train OpenFlamingo, please ensure your environment matches that of `environment.yml`.
+
+## Data
+Our codebase uses [WebDataset](https://github.com/webdataset/webdataset) to efficiently load `.tar` files containing image and text sequences. We recommend resampling shards with replacement during training using the `--dataset_resampled` flag. 
+
+### LAION-2B Dataset
+[LAION-2B](https://arxiv.org/abs/2210.08402) contains 2B web-scraped (image, text) pairs. 
+We use [img2dataset](https://github.com/rom1504/img2dataset) to download this dataset into tar files.
+
+### Multimodal C4 Dataset
+We train on the full version of [Multimodal C4 (MMC4)](https://github.com/allenai/mmc4), which includes 103M documents of web-scraped, interleaved image-text sequences. During training, we truncate sequences to 256 text tokens and six images per sequence.
+
+Our codebase expects `.tar` files containing `.json` files, which include raw images encoded in base64.
+We provide scripts to convert MMC4 to this format: 
+
+1. Download the MMC4 shards into `.zip` files using [the MMC4-provided scripts](https://github.com/allenai/mmc4/tree/main/scripts) (e.g., `fewer_facesv2.sh`).
+2. Download the MMC4 raw images into an image directory using [the MMC4-provided scripts](https://github.com/allenai/mmc4/tree/main/scripts) (e.g., `download_images.py`).
+2. Run `scripts/convert_mmc4_to_wds.py` to convert the downloaded items into the expected tar files.
+
+### ChatGPT-generated sequences
+A subset of our models (listed below) were also trained on experimental ChatGPT-generated (image, text) sequences, where images are pulled from LAION. The shards containing these sequences can be found at [this CodaLab worksheet](https://worksheets.codalab.org/worksheets/0xdcd888ff7c754ae680c5e038f6ed1d9b). We are unable to distribute raw images in the released shards; images must be pre-downloaded from the urls in the json files and converted to base64 before using this data for training in our codebase.
+
+Models trained with ChatGPT-generated sequences:
+
+* OpenFlamingo-4B-vitl-rpj3b
+* OpenFlamingo-4B-vitl-rpj3b-langinstruct
+
+## Example training command
+We provide a sample Slurm training script in `scripts/`. You can also modify the following command:
+
+```
+torchrun --nnodes=1 --nproc_per_node=4 train.py \
+  --lm_path anas-awadalla/mpt-1b-redpajama-200b \
+  --tokenizer_path anas-awadalla/mpt-1b-redpajama-200b \
+  --cross_attn_every_n_layers 1 \
+  --dataset_resampled \
+  --batch_size_mmc4 32 \
+  --batch_size_laion 64 \
+  --train_num_samples_mmc4 125000\
+  --train_num_samples_laion 250000 \
+  --loss_multiplier_laion 0.2 \
+  --workers=4 \
+  --run_name OpenFlamingo-3B-vitl-mpt1b \
+  --num_epochs 480 \
+  --warmup_steps  1875 \
+  --mmc4_textsim_threshold 0.24 \
+  --laion_shards "/path/to/shards/shard-{0000..0999}.tar" \
+  --mmc4_shards "/path/to/shards/shard-{0000..0999}.tar" \
+  --report_to_wandb
+```
+*Note: The MPT-1B [base](https://huggingface.co/mosaicml/mpt-1b-redpajama-200b)  and [instruct](https://huggingface.co/mosaicml/mpt-1b-redpajama-200b-dolly) modeling code does not accept the `labels` kwarg or compute cross-entropy loss directly within `forward()`, as expected by our codebase. We suggest using a modified version of the MPT-1B models found [here](https://huggingface.co/anas-awadalla/mpt-1b-redpajama-200b) and [here](https://huggingface.co/anas-awadalla/mpt-1b-redpajama-200b-dolly).*
+
+## Distributed training
+
+By default, `train.py` uses Pytorch's [DistributedDataParallel](https://pytorch.org/docs/stable/torch.nn.parallel.DistributedDataParallel.html) for training. 
+To use [FullyShardedDataParallel](https://pytorch.org/docs/stable/fsdp.html), use the `--fsdp` flag. 
+
+Some notes on FSDP:
+
+* We recommend using the `--fsdp_use_orig_params` flag. If `--fsdp` is on without this flag, all language model embeddings will be unfrozen during training. (In contrast, the default behavior is to only train the newly added `<image>` and `<|endofchunk|>` tokens.)
+    * Note: we've encountered issues using OPT with this flag. Other language models should be compatible.
+* Our current FSDP wrapping strategy does not permit training language model embeddings that use tied weights (i.e., tied input / output embeddings). To train such models with FSDP, the language model embeddings must be frozen with the `--freeze_lm_embeddings` flag.
+
+We also implement gradient checkpointing and mixed precision training. Use the `--gradient_checkpointing` and `--precision` arguments respectively.

open_flamingo/train/__init__.py

@@ -0,0 +1 @@
+

open_flamingo/train/data.py

@@ -0,0 +1,492 @@
+"""
+Preprocess and load datasets for training.
+"""
+
+import functools
+import io
+import json
+import math
+import re
+import random
+import numpy as np
+import torch
+import torchvision
+import webdataset as wds
+from PIL import Image
+import base64
+from scipy.optimize import linear_sum_assignment
+
+from data_utils import *
+
+Image.MAX_IMAGE_PIXELS = 1000000000
+N_CHANNELS = 3
+MIN_KB = 10
+_SHARD_SHUFFLE_SIZE = 2000
+_SHARD_SHUFFLE_INITIAL = 500
+_SAMPLE_SHUFFLE_SIZE = 5000
+_SAMPLE_SHUFFLE_INITIAL = 1000
+
+try:
+    import horovod.torch as hvd
+except ImportError:
+    hvd = None
+
+
+def preprocess_image(sample, image_processor):
+    """
+    Convert images to tensors for training.
+    Augmentations: random horizontal flip.
+    Normalization handled by wds.
+    """
+    image = [image_processor(s).unsqueeze(0) for s in sample]
+    image = torch.cat(image, dim=0)
+    image = torchvision.transforms.RandomHorizontalFlip(p=0.5)(image)
+    return image
+
+
+def filter_no_caption_or_no_image(sample):
+    """
+    Filter out LAION samples with no caption or no image.
+    """
+    return ("txt" in sample) and (
+        "png" in sample or "jpg" in sample or "jpeg" in sample
+    )
+
+
+def preprocess_laion_text(sample, tokenizer, max_tokens=32):
+    """
+    Preprocess text for LAION.
+    Captions are truncated to 32 tokens by default.
+    """
+    tokenizer.padding_side = "right"
+    sample = [
+        (f"<image>{s.strip()}<|endofchunk|>{tokenizer.eos_token}") for s in sample
+    ]
+    text = tokenizer(
+        sample,
+        max_length=max_tokens,
+        padding="longest",
+        truncation="only_first",
+        return_tensors="pt",
+    )
+    return text["input_ids"], text["attention_mask"]
+
+
+def preprocess_gpt_interleaved(
+    info, tokenizer, clip_processor, min_num_images, max_num_images, max_tokens=256
+):
+    """
+    Preprocess a ChatGPT-generated image-text sequence.
+    """
+    text = info["example"]
+    text = re.sub(r"_!_IMAGE\d+_!_", "<|endofchunk|><image>", text)
+
+    # convert images from base64 to PIL
+    images = []
+    for image_key in range(1, len(info["image_map"]) + 1):
+        image_base64 = info["image_map"][f"_!_IMAGE{image_key}_!_"]["base64_image"]
+        rawbytes = base64.b64decode(image_base64)
+        images.append(Image.open(io.BytesIO(rawbytes)).convert("RGB"))
+
+    # preprocess and pad images
+    images_tensors = preprocess_image(images, clip_processor)
+    keep_ixs = range(min(len(images_tensors), max_num_images))
+    images_tensors = images_tensors[keep_ixs]
+    if len(images_tensors) < max_num_images:
+        zero_padding = torch.zeros(
+            (max_num_images - len(images_tensors), 3, 224, 224), dtype=torch.float
+        )
+        images_tensors = torch.cat((images_tensors, zero_padding), dim=0)
+
+    # preprocess and tokenize text
+    text = text.replace("<|endofchunk|>", "", 1)  # but remove first eoc
+    # whitespace cleanup
+    text = (
+        text.replace(" <|endofchunk|>", "<|endofchunk|>")
+        .replace("<image> ", "<image>")
+        .replace(" <image>", "<image>")
+    )
+
+    indices = [m.start() for m in re.finditer("<image>", text)]
+    if len(indices) > max_num_images:
+        start_index = indices[max_num_images - 1]
+        text = text[:start_index]
+
+    text = f"{text}<|endofchunk|>{tokenizer.eos_token}"
+    tokenizer.padding_side = "right"
+    text_tensor = tokenizer(
+        text,
+        max_length=max_tokens,
+        truncation=True,
+        padding="max_length",
+        return_tensors="pt",
+    )
+
+    # reject sequences with too few images after truncation
+    num_images = torch.count_nonzero(
+        text_tensor["input_ids"]
+        == tokenizer.additional_special_tokens_ids[
+            tokenizer.additional_special_tokens.index("<image>")
+        ]
+    )
+    if num_images < min_num_images:
+        raise ValueError(f"Fewer than {min_num_images} images in sample")
+
+    return (images_tensors, (text_tensor["input_ids"], text_tensor["attention_mask"]))
+
+
+def preprocess_interleaved(
+    sample,
+    tokenizer,
+    clip_processor,
+    sim_threshold,
+    min_num_images,
+    max_num_images,
+    max_tokens=256,
+):
+    """
+    Preprocess an interleaved image-text sequence, either by calling preprocess_gpt_interleaved (if the sequence
+    is ChatGPT-generated) or by preprocessing in this function (if the sequences is from MMC4).
+    """
+    info = json.loads(sample[0])
+    if "is_gpt" in info:
+        return preprocess_gpt_interleaved(
+            info, tokenizer, clip_processor, min_num_images, max_num_images, max_tokens
+        )
+
+    sentences = info["text_list"]
+    sim_matrix = info["similarity_matrix"]
+
+    # load images first to find which ones are valid
+    valid_images, valid_image_indices = [], []
+    for i, sample_image in enumerate(info["image_info"]):
+        if "image_base64" not in sample_image:
+            continue
+        image_base64 = sample_image["image_base64"]
+        rawbytes = base64.b64decode(image_base64)
+
+        # filter to images >= 10KB
+        if len(rawbytes) // 1000 <= MIN_KB:
+            continue
+
+        image = Image.open(io.BytesIO(rawbytes)).convert("RGB")
+        valid_images.append(image)
+        valid_image_indices.append(i)
+
+    if len(valid_image_indices) == 0:
+        raise ValueError("No images in sample")
+
+    sim_matrix = np.array(sim_matrix)  # of shape images x sentences
+    sim_matrix = sim_matrix[valid_image_indices]
+
+    # negate the similarities to turn then into costs
+    cost_matrix = -sim_matrix
+    # find one to one assignements
+    image_indices, sentence_indices = linear_sum_assignment(cost_matrix)
+
+    images, sentence_ixs = [], []
+    for i, sim_ix in zip(image_indices, sentence_indices):
+        sim_score = sim_matrix[i][sim_ix]
+
+        if sim_score < sim_threshold:
+            continue
+
+        images.append(valid_images[i])
+        sentence_ixs.append(sim_ix)
+
+    if len(images) == 0:
+        raise ValueError("No images in sample")
+
+    # preprocess and pad images
+    images_tensors = preprocess_image(images, clip_processor)
+    keep_ixs = range(min(len(images_tensors), max_num_images))
+    images_tensors = images_tensors[keep_ixs]
+    sentence_ixs = [sentence_ixs[ix] for ix in keep_ixs]
+    if len(images_tensors) < max_num_images:
+        zero_padding = torch.zeros(
+            (
+                max_num_images - len(images_tensors),
+                N_CHANNELS,
+                images_tensors[0].shape[1],
+                images_tensors[0].shape[2],
+            ),
+            dtype=torch.float,
+        )
+        images_tensors = torch.cat((images_tensors, zero_padding), dim=0)
+
+    # preprocess and tokenize text
+    # add in <image> and <eoc> tokens
+    for ix in sentence_ixs:
+        sentences[ix] = f"<|endofchunk|><image>{sentences[ix]}"
+    text = " ".join(sentences)
+    text = text.replace("<|endofchunk|>", "", 1)  # but remove first eoc
+    # whitespace cleanup
+    text = (
+        text.replace(" <|endofchunk|>", "<|endofchunk|>")
+        .replace("<image> ", "<image>")
+        .replace(" <image>", "<image>")
+    )
+    text = f"{text}<|endofchunk|>{tokenizer.eos_token}"
+    tokenizer.padding_side = "right"
+    text_tensor = tokenizer(
+        text,
+        max_length=max_tokens,
+        truncation=True,
+        padding="max_length",
+        return_tensors="pt",
+    )
+
+    # reject sequences with too few images (after truncation)
+    num_images = torch.count_nonzero(
+        text_tensor["input_ids"]
+        == tokenizer.additional_special_tokens_ids[
+            tokenizer.additional_special_tokens.index("<image>")
+        ]
+    )
+    if num_images < min_num_images:
+        raise ValueError(f"Fewer than {min_num_images} images in sample")
+    elif (
+        num_images == 1 and random.random() <= 0.5
+    ):  # 50% chance of keeping single image samples
+        raise ValueError("Only one image in sample")
+
+    # avoid the situation where there's one <image> token and it's at the end
+    if (
+        num_images == 1
+        and text_tensor["input_ids"][:, -1]
+        == tokenizer.additional_special_tokens_ids[
+            tokenizer.additional_special_tokens.index("<image>")
+        ]
+    ):
+        raise ValueError(
+            "Only one image at the end of sample, so labels will all be -100"
+        )
+
+    return (
+        images_tensors,
+        (text_tensor["input_ids"], text_tensor["attention_mask"]),
+    )
+
+
+def get_mmc4_dataset(args, image_processor, tokenizer, epoch=0, floor=False):
+    """
+    Initialize webdataset for MMC4 / ChatGPT sequences
+    """
+    input_shards = args.mmc4_shards
+    assert input_shards is not None
+    resampled = getattr(args, "dataset_resampled", False)
+
+    num_samples, num_shards = get_dataset_size(input_shards)
+    num_samples = None
+    if not num_samples:
+        num_samples = args.train_num_samples_mmc4
+        if not num_samples:
+            raise RuntimeError(
+                "Currently, number of dataset samples must be specified for training dataset. "
+                "Please specify via `--train-num-samples` if no dataset length info present."
+            )
+
+    # create a shared epoch store to sync epoch to dataloader worker proc
+    shared_epoch = SharedEpoch(epoch=epoch)
+    if resampled:
+        pipeline = [
+            ResampledShards2(input_shards, deterministic=True, epoch=shared_epoch)
+        ]
+    else:
+        pipeline = [wds.SimpleShardList(input_shards)]
+
+    preprocess_fn = functools.partial(
+        preprocess_interleaved,
+        clip_processor=image_processor,
+        tokenizer=tokenizer,
+        sim_threshold=args.mmc4_textsim_threshold,
+        min_num_images=args.mmc4_min_num_images,
+        max_num_images=args.mmc4_max_num_images,
+    )
+
+    # at this point we have an iterator over all the shards
+    if not resampled:
+        pipeline.extend(
+            [
+                detshuffle2(
+                    bufsize=_SHARD_SHUFFLE_SIZE,
+                    initial=_SHARD_SHUFFLE_INITIAL,
+                    seed=args.seed,
+                    epoch=shared_epoch,
+                ),
+                wds.split_by_node,
+                wds.split_by_worker,
+            ]
+        )
+    pipeline.extend(
+        [
+            # at this point, we have an iterator over the shards assigned to each worker at each node
+            # wds.tarfile_to_samples(handler=log_and_continue),
+            tarfile_to_samples_nothrow,
+            wds.shuffle(
+                bufsize=_SAMPLE_SHUFFLE_SIZE,
+                initial=_SAMPLE_SHUFFLE_INITIAL,
+            ),
+        ]
+    )
+
+    pipeline.extend(
+        [
+            wds.to_tuple("json", handler=log_and_continue),
+            wds.map(preprocess_fn, handler=log_and_continue),
+            wds.batched(args.batch_size_mmc4, partial=False),
+        ]
+    )
+
+    dataset = wds.DataPipeline(*pipeline)
+    if not resampled:
+        assert (
+            num_shards >= args.workers * args.world_size
+        ), "number of shards must be >= total workers"
+    # roll over and repeat a few samples to get same number of full batches on each node
+    round_fn = math.floor if floor else math.ceil
+    global_batch_size = args.batch_size_mmc4 * args.world_size
+    num_batches = round_fn(num_samples / global_batch_size)
+    num_workers = max(1, args.workers)
+    num_worker_batches = round_fn(num_batches / num_workers)  # per dataloader worker
+    num_batches = num_worker_batches * num_workers
+    num_samples = num_batches * global_batch_size
+    # each worker is iterating over this
+    dataset = dataset.with_epoch(num_worker_batches)
+
+    dataloader = wds.WebLoader(
+        dataset,
+        batch_size=None,
+        shuffle=False,
+        num_workers=args.workers,
+        persistent_workers=True,
+    )
+
+    # add meta-data to dataloader instance for convenience
+    dataloader.num_batches = num_batches
+    dataloader.num_samples = num_samples
+
+    return DataInfo(dataloader=dataloader, shared_epoch=shared_epoch)
+
+
+def get_laion_dataset(args, image_processor, tokenizer, epoch=0, floor=False):
+    """
+    Initialize webdataset for LAION data
+    """
+    input_shards = args.laion_shards
+    assert input_shards is not None
+    resampled = getattr(args, "dataset_resampled", False)
+
+    num_samples, num_shards = get_dataset_size(input_shards)
+    num_samples = None
+    if not num_samples:
+        num_samples = args.train_num_samples_laion
+        if not num_samples:
+            raise RuntimeError(
+                "Currently, number of dataset samples must be specified for training dataset. "
+                "Please specify via `--train-num-samples` if no dataset length info present."
+            )
+
+    # create a shared epoch store to sync epoch to dataloader worker proc
+    shared_epoch = SharedEpoch(epoch=epoch)
+    if resampled:
+        pipeline = [
+            ResampledShards2(input_shards, deterministic=True, epoch=shared_epoch)
+        ]
+    else:
+        pipeline = [wds.SimpleShardList(input_shards)]
+
+    # create two preprocess functions that take in the passed in image_processor and tokenizer
+    preprocess_image_fn = functools.partial(
+        preprocess_image, image_processor=image_processor
+    )
+    preprocess_text_fn = functools.partial(preprocess_laion_text, tokenizer=tokenizer)
+
+    # at this point we have an iterator over all the shards
+    if not resampled:
+        pipeline.extend(
+            [
+                detshuffle2(
+                    bufsize=_SHARD_SHUFFLE_SIZE,
+                    initial=_SHARD_SHUFFLE_INITIAL,
+                    seed=args.seed,
+                    epoch=shared_epoch,
+                ),
+                wds.split_by_node,
+                wds.split_by_worker,
+            ]
+        )
+    pipeline.extend(
+        [
+            # at this point, we have an iterator over the shards assigned to each worker at each node
+            # wds.tarfile_to_samples(handler=log_and_continue),
+            tarfile_to_samples_nothrow,
+            wds.shuffle(
+                bufsize=_SAMPLE_SHUFFLE_SIZE,
+                initial=_SAMPLE_SHUFFLE_INITIAL,
+            ),
+        ]
+    )
+
+    pipeline.extend(
+        [
+            wds.select(filter_no_caption_or_no_image),
+            wds.decode("pilrgb", handler=log_and_continue),
+            wds.to_tuple("jpg;png;jpeg", "txt", handler=log_and_continue),
+            wds.batched(args.batch_size_laion, partial=False),
+            wds.map_tuple(
+                preprocess_image_fn, preprocess_text_fn, handler=log_and_continue
+            ),
+        ]
+    )
+
+    dataset = wds.DataPipeline(*pipeline)
+    if not resampled:
+        assert (
+            num_shards >= args.workers * args.world_size
+        ), "number of shards must be >= total workers"
+    # roll over and repeat a few samples to get same number of full batches on each node
+    round_fn = math.floor if floor else math.ceil
+    global_batch_size = args.batch_size_laion * args.world_size
+    num_batches = round_fn(num_samples / global_batch_size)
+    num_workers = max(1, args.workers)
+    num_worker_batches = round_fn(num_batches / num_workers)  # per dataloader worker
+    num_batches = num_worker_batches * num_workers
+    num_samples = num_batches * global_batch_size
+    # each worker is iterating over this
+    dataset = dataset.with_epoch(num_worker_batches)
+
+    dataloader = wds.WebLoader(
+        dataset,
+        batch_size=None,
+        shuffle=False,
+        num_workers=args.workers,
+        persistent_workers=True,
+    )
+
+    # add meta-data to dataloader instance for convenience
+    dataloader.num_batches = num_batches
+    dataloader.num_samples = num_samples
+
+    return DataInfo(dataloader=dataloader, shared_epoch=shared_epoch)
+
+
+def get_dataset_fn(dataset_type):
+    """
+    Helper function to get the dataset function based on the dataset type
+    """
+    if dataset_type == "image_text":
+        return get_laion_dataset
+    elif dataset_type == "mmc4":
+        return get_mmc4_dataset
+    else:
+        raise ValueError(f"Unsupported dataset type: {dataset_type}")
+
+
+def get_data(args, image_processor, tokenizer, dataset_type, epoch=0):
+    """
+    Interface for getting the webdatasets
+    """
+    return get_dataset_fn(dataset_type)(
+        args, image_processor=image_processor, epoch=epoch, tokenizer=tokenizer
+    )

open_flamingo/train/data_utils.py

@@ -0,0 +1,234 @@
+"""
+Util functions for initializing webdataset objects
+"""
+
+import ast
+import json
+import logging
+import os
+import random
+import sys
+from dataclasses import dataclass
+from multiprocessing import Value
+
+import braceexpand
+import numpy as np
+import webdataset as wds
+from PIL import Image
+from torch.utils.data import DataLoader, IterableDataset, get_worker_info
+from torch.utils.data.distributed import DistributedSampler
+from webdataset.filters import _shuffle
+from webdataset.tariterators import (
+    base_plus_ext,
+    tar_file_expander,
+    url_opener,
+    valid_sample,
+)
+
+try:
+    import horovod.torch as hvd
+except ImportError:
+    hvd = None
+
+
+class SharedEpoch:
+    def __init__(self, epoch: int = 0):
+        self.shared_epoch = Value("i", epoch)
+
+    def set_value(self, epoch):
+        self.shared_epoch.value = epoch
+
+    def get_value(self):
+        return self.shared_epoch.value
+
+
+@dataclass
+class DataInfo:
+    dataloader: DataLoader
+    sampler: DistributedSampler = None
+    shared_epoch: SharedEpoch = None
+
+    def set_epoch(self, epoch):
+        if self.shared_epoch is not None:
+            self.shared_epoch.set_value(epoch)
+        if self.sampler is not None and isinstance(self.sampler, DistributedSampler):
+            self.sampler.set_epoch(epoch)
+
+
+def get_dataset_size(shards):
+    shards_list = list(braceexpand.braceexpand(shards))
+    dir_path = os.path.dirname(shards[0])
+    sizes_filename = os.path.join(dir_path, "sizes.json")
+    len_filename = os.path.join(dir_path, "__len__")
+    if os.path.exists(sizes_filename):
+        sizes = json.load(open(sizes_filename, "r"))
+        total_size = sum(
+            [
+                int(sizes[os.path.basename(shard)])
+                if os.path.basename(shard) in sizes
+                else 0
+                for shard in shards_list
+            ]
+        )
+    elif os.path.exists(len_filename):
+        # FIXME this used to be eval(open(...)) but that seemed rather unsafe
+        total_size = ast.literal_eval(open(len_filename, "r").read())
+    else:
+        total_size = None  # num samples undefined
+        # some common dataset sizes (at time of authors last download)
+        # CC3M (train): 2905954
+        # CC12M: 10968539
+        # LAION-400M: 407332084
+        # LAION-2B (english): 2170337258
+    num_shards = len(shards_list)
+    return total_size, num_shards
+
+
+def count_samples(dataloader):
+    os.environ["WDS_EPOCH"] = "0"
+    n_elements, n_batches = 0, 0
+    for images, texts in dataloader:
+        n_batches += 1
+        n_elements += len(images)
+        assert len(images) == len(texts)
+    return n_elements, n_batches
+
+
+def log_and_continue(exn):
+    """Call in an exception handler to ignore any exception, issue a warning, and continue."""
+    logging.warning(f"Handling webdataset error ({repr(exn)}). Ignoring.")
+    return True
+
+
+def group_by_keys_nothrow(
+    data, keys=base_plus_ext, lcase=True, suffixes=None, handler=None
+):
+    """Return function over iterator that groups key, value pairs into samples.
+
+    :param keys: function that splits the key into key and extension (base_plus_ext)
+    :param lcase: convert suffixes to lower case (Default value = True)
+    """
+    current_sample = None
+    for filesample in data:
+        assert isinstance(filesample, dict)
+        fname, value = filesample["fname"], filesample["data"]
+        prefix, suffix = keys(fname)
+        if prefix is None:
+            continue
+        if lcase:
+            suffix = suffix.lower()
+        # FIXME webdataset version throws if suffix in current_sample, but we have a potential for
+        #  this happening in the current LAION400m dataset if a tar ends with same prefix as the next
+        #  begins, rare, but can happen since prefix aren't unique across tar files in that dataset
+        if (
+            current_sample is None
+            or prefix != current_sample["__key__"]
+            or suffix in current_sample
+        ):
+            if valid_sample(current_sample):
+                yield current_sample
+            current_sample = dict(__key__=prefix, __url__=filesample["__url__"])
+        if suffixes is None or suffix in suffixes:
+            current_sample[suffix] = value
+    if valid_sample(current_sample):
+        yield current_sample
+
+
+def tarfile_to_samples_nothrow(src, handler=log_and_continue):
+    # NOTE this is a re-impl of the webdataset impl with group_by_keys that doesn't throw
+    streams = url_opener(src, handler=handler)
+    files = tar_file_expander(streams, handler=handler)
+    samples = group_by_keys_nothrow(files, handler=handler)
+    return samples
+
+
+def pytorch_worker_seed(increment=0):
+    """get dataloader worker seed from pytorch"""
+    worker_info = get_worker_info()
+    if worker_info is not None:
+        # favour using the seed already created for pytorch dataloader workers if it exists
+        seed = worker_info.seed
+        if increment:
+            # space out seed increments so they can't overlap across workers in different iterations
+            seed += increment * max(1, worker_info.num_workers)
+        return seed
+    # fallback to wds rank based seed
+    return wds.utils.pytorch_worker_seed()
+
+
+class detshuffle2(wds.PipelineStage):
+    def __init__(
+        self,
+        bufsize=1000,
+        initial=100,
+        seed=0,
+        epoch=-1,
+    ):
+        self.bufsize = bufsize
+        self.initial = initial
+        self.seed = seed
+        self.epoch = epoch
+
+    def run(self, src):
+        if isinstance(self.epoch, SharedEpoch):
+            epoch = self.epoch.get_value()
+        else:
+            # NOTE: this is epoch tracking is problematic in a multiprocess (dataloader workers or train)
+            # situation as different workers may wrap at different times (or not at all).
+            self.epoch += 1
+            epoch = self.epoch
+        rng = random.Random()
+        if self.seed < 0:
+            # If seed is negative, we use the worker's seed, this will be different across all nodes/workers
+            seed = pytorch_worker_seed(epoch)
+        else:
+            # This seed to be deterministic AND the same across all nodes/workers in each epoch
+            seed = self.seed + epoch
+        rng.seed(seed)
+        return _shuffle(src, self.bufsize, self.initial, rng)
+
+
+class ResampledShards2(IterableDataset):
+    """An iterable dataset yielding a list of urls."""
+
+    def __init__(
+        self,
+        urls,
+        nshards=sys.maxsize,
+        worker_seed=None,
+        deterministic=False,
+        epoch=-1,
+    ):
+        """Sample shards from the shard list with replacement.
+        :param urls: a list of URLs as a Python list or brace notation string
+        """
+        super().__init__()
+        urls = wds.shardlists.expand_urls(urls)
+        self.urls = urls
+        assert isinstance(self.urls[0], str)
+        self.nshards = nshards
+        self.rng = random.Random()
+        self.worker_seed = worker_seed
+        self.deterministic = deterministic
+        self.epoch = epoch
+
+    def __iter__(self):
+        """Return an iterator over the shards."""
+        if isinstance(self.epoch, SharedEpoch):
+            epoch = self.epoch.get_value()
+        else:
+            # NOTE: this is epoch tracking is problematic in a multiprocess (dataloader workers or train)
+            # situation as different workers may wrap at different times (or not at all).
+            self.epoch += 1
+            epoch = self.epoch
+
+        if self.deterministic:
+            # reset seed w/ epoch if deterministic
+            if self.worker_seed is None:
+                # pytorch worker seed should be deterministic due to being init by arg.seed + rank + worker id
+                seed = pytorch_worker_seed(epoch)
+            else:
+                seed = self.worker_seed() + epoch
+            self.rng.seed(seed)
+        for _ in range(self.nshards):
+            yield dict(url=self.rng.choice(self.urls))

open_flamingo/train/distributed.py

@@ -0,0 +1,132 @@
+"""
+Util functions for setting up distributed training.
+Credit: https://github.com/mlfoundations/open_clip/blob/main/src/training/distributed.py
+"""
+
+import os
+import torch
+
+try:
+    import horovod.torch as hvd
+except ImportError:
+    hvd = None
+
+
+def is_global_master(args):
+    return args.rank == 0
+
+
+def is_local_master(args):
+    return args.local_rank == 0
+
+
+def is_master(args, local=False):
+    return is_local_master(args) if local else is_global_master(args)
+
+
+def is_using_horovod():
+    # NOTE w/ horovod run, OMPI vars should be set, but w/ SLURM PMI vars will be set
+    # Differentiating between horovod and DDP use via SLURM may not be possible, so horovod arg still required...
+    ompi_vars = ["OMPI_COMM_WORLD_RANK", "OMPI_COMM_WORLD_SIZE"]
+    pmi_vars = ["PMI_RANK", "PMI_SIZE"]
+    if all([var in os.environ for var in ompi_vars]) or all(
+        [var in os.environ for var in pmi_vars]
+    ):
+        return True
+    else:
+        return False
+
+
+def is_using_distributed():
+    if "WORLD_SIZE" in os.environ:
+        return int(os.environ["WORLD_SIZE"]) > 1
+    if "SLURM_NTASKS" in os.environ:
+        return int(os.environ["SLURM_NTASKS"]) > 1
+    return False
+
+
+def world_info_from_env():
+    local_rank = 0
+    for v in (
+        "LOCAL_RANK",
+        "MPI_LOCALRANKID",
+        "SLURM_LOCALID",
+        "OMPI_COMM_WORLD_LOCAL_RANK",
+    ):
+        if v in os.environ:
+            local_rank = int(os.environ[v])
+            break
+    global_rank = 0
+    for v in ("RANK", "PMI_RANK", "SLURM_PROCID", "OMPI_COMM_WORLD_RANK"):
+        if v in os.environ:
+            global_rank = int(os.environ[v])
+            break
+    world_size = 1
+    for v in ("WORLD_SIZE", "PMI_SIZE", "SLURM_NTASKS", "OMPI_COMM_WORLD_SIZE"):
+        if v in os.environ:
+            world_size = int(os.environ[v])
+            break
+
+    return local_rank, global_rank, world_size
+
+
+def init_distributed_device(args):
+    # Distributed training = training on more than one GPU.
+    # Works in both single and multi-node scenarios.
+    args.distributed = False
+    args.world_size = 1
+    args.rank = 0  # global rank
+    args.local_rank = 0
+    if args.horovod:
+        assert hvd is not None, "Horovod is not installed"
+        hvd.init()
+        args.local_rank = int(hvd.local_rank())
+        args.rank = hvd.rank()
+        args.world_size = hvd.size()
+        args.distributed = True
+        os.environ["LOCAL_RANK"] = str(args.local_rank)
+        os.environ["RANK"] = str(args.rank)
+        os.environ["WORLD_SIZE"] = str(args.world_size)
+    elif is_using_distributed():
+        if "SLURM_PROCID" in os.environ:
+            # DDP via SLURM
+            args.local_rank, args.rank, args.world_size = world_info_from_env()
+            # SLURM var -> torch.distributed vars in case needed
+            os.environ["LOCAL_RANK"] = str(args.local_rank)
+            os.environ["RANK"] = str(args.rank)
+            os.environ["WORLD_SIZE"] = str(args.world_size)
+            torch.distributed.init_process_group(
+                backend=args.dist_backend,
+                init_method=args.dist_url,
+                world_size=args.world_size,
+                rank=args.rank,
+            )
+        else:
+            # DDP via torchrun, torch.distributed.launch
+            args.local_rank, _, _ = world_info_from_env()
+            torch.distributed.init_process_group(
+                backend=args.dist_backend, init_method=args.dist_url
+            )
+            args.world_size = torch.distributed.get_world_size()
+            args.rank = torch.distributed.get_rank()
+        args.distributed = True
+    else:
+        # needed to run on single gpu
+        torch.distributed.init_process_group(
+            backend=args.dist_backend,
+            init_method=args.dist_url,
+            world_size=1,
+            rank=0,
+        )
+
+    if torch.cuda.is_available():
+        if args.distributed and not args.no_set_device_rank:
+            device = "cuda:%d" % args.local_rank
+        else:
+            device = "cuda:0"
+        torch.cuda.set_device(device)
+    else:
+        device = "cpu"
+    args.device = device
+    device = torch.device(device)
+    return device

open_flamingo/train/train.py

@@ -0,0 +1,484 @@
+""" Main training script """
+
+import argparse
+import glob
+import os
+import random
+
+import numpy as np
+import torch
+import wandb
+from data import get_data
+from distributed import init_distributed_device, world_info_from_env
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+from train_utils import (
+    train_one_epoch,
+    get_mp_policy_dtype,
+    save_checkpoint,
+)
+from transformers import (
+    get_constant_schedule_with_warmup,
+    get_cosine_schedule_with_warmup,
+    get_linear_schedule_with_warmup,
+)
+
+from torch.distributed.fsdp import (
+    CPUOffload,
+    MixedPrecision,
+    ShardingStrategy,
+    BackwardPrefetch,
+)
+from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
+    checkpoint_wrapper,
+    CheckpointWrapper,
+    CheckpointImpl,
+    apply_activation_checkpointing,
+)
+from torch.distributed.fsdp._init_utils import _init_intra_and_inter_node_groups
+from torch.distributed.distributed_c10d import _get_default_group
+import functools
+
+from open_flamingo import create_model_and_transforms
+
+
+def random_seed(seed=42, rank=0):
+    torch.manual_seed(seed + rank)
+    np.random.seed(seed + rank)
+    random.seed(seed + rank)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    # model configuration args
+    parser.add_argument("--vision_encoder_path", default="ViT-L-14", type=str)
+    parser.add_argument("--vision_encoder_pretrained", default="openai", type=str)
+    parser.add_argument("--lm_path", default="facebook/opt-1.3b", type=str)
+    parser.add_argument(
+        "--tokenizer_path",
+        default="facebook/opt-30b",
+        type=str,
+        help="path to tokenizer",
+    )
+    parser.add_argument(
+        "--cross_attn_every_n_layers",
+        type=int,
+        default=1,
+        help="how often to add a cross-attention layer after each transformer layer",
+    )
+
+    # training args
+    parser.add_argument(
+        "--run_name",
+        type=str,
+        default="openflamingo3B",
+        help="used to name saving directory and wandb run",
+    )
+    parser.add_argument(
+        "--resume_from_checkpoint",
+        type=str,
+        help="path to checkpoint to resume from, this should contain model, optimizer, and lr_scheduler states. if there exists a checkpoint in the dir named run_name, we will resume from that checkpoint by default",
+        default=None,
+    )
+    parser.add_argument(
+        "--delete_previous_checkpoint",
+        action="store_true",
+        help="delete previous checkpoint when saving new checkpoint",
+    )
+    parser.add_argument("--batch_size_mmc4", type=int, default=128)
+    parser.add_argument("--batch_size_laion", type=int, default=128)
+    parser.add_argument("--gradient_accumulation_steps", type=int, default=1)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--learning_rate", default=1e-4, type=float)
+    parser.add_argument(
+        "--lr_scheduler",
+        default="constant",
+        type=str,
+        help="constant, linear, or cosine",
+    )
+    parser.add_argument("--loss_multiplier_mmc4", type=float, default=1.0)
+    parser.add_argument("--loss_multiplier_laion", type=float, default=1.0)
+    parser.add_argument("--warmup_steps", default=5000, type=int)
+    parser.add_argument("--weight_decay", default=0.1, type=float)
+    parser.add_argument(
+        "--precision",
+        choices=["amp_bf16", "amp_bfloat16", "bf16", "fp16", "fp32"],
+        default="fp32",
+        help="Floating point precision.",
+    )
+    parser.add_argument(
+        "--gradient_checkpointing",
+        action="store_true",
+        help="whether to train with gradient/activation checkpointing",
+    )
+    parser.add_argument(
+        "--num_epochs",
+        type=int,
+        default=1,
+        help="we define an 'epoch' as a fixed number of examples (train_num_samples_mmc4, train_num_samples_laion), not a pass through the entire dataset",
+    )
+    parser.add_argument("--offline", action="store_true")
+    parser.add_argument(
+        "--freeze_lm_embeddings",
+        action="store_true",
+        help="if True, we freeze the LM embeddings during training. Otherwise, we train the <image> and <|endofchunk|> embeddings.",
+    )
+    parser.add_argument(
+        "--logging_steps", type=int, default=100, help="log loss every n steps"
+    )
+
+    # data args
+    parser.add_argument(
+        "--laion_shards",
+        type=str,
+        help="path to laion shards, this should be a glob pattern such as /path/to/shards/shard-{0000..0999}.tar",
+    )
+    parser.add_argument(
+        "--mmc4_shards",
+        type=str,
+        help="path to c4 shards, this should be a glob pattern such as /path/to/shards/shard-{0000..0999}.tar",
+    )
+    parser.add_argument("--workers", type=int, default=1)
+    parser.add_argument("--train_num_samples_mmc4", type=int, default=10000)
+    parser.add_argument("--train_num_samples_laion", type=int, default=10000)
+    parser.add_argument("--dataset_resampled", action="store_true")
+    parser.add_argument(
+        "--mmc4_textsim_threshold",
+        default=30,
+        type=float,
+        help="threshold for filtering images in mmc4 based on image-text similarity",
+    )
+    parser.add_argument(
+        "--mmc4_max_num_images",
+        default=6,
+        type=int,
+        help="max number of images per sequence in mmc4 / chatgpt",
+    )
+    parser.add_argument(
+        "--mmc4_min_num_images",
+        default=1,
+        type=int,
+        help="min number of images per sequence in mmc4 / chatgpt",
+    )
+
+    # distributed training args
+    parser.add_argument(
+        "--dist-url",
+        default="env://",
+        type=str,
+        help="url used to set up distributed training",
+    )
+    parser.add_argument(
+        "--dist-backend", default="nccl", type=str, help="distributed backend"
+    )
+    parser.add_argument(
+        "--horovod",
+        default=False,
+        action="store_true",
+        help="Use horovod for distributed training.",
+    )
+    parser.add_argument(
+        "--no-set-device-rank",
+        default=False,
+        action="store_true",
+        help="Don't set device index from local rank (when CUDA_VISIBLE_DEVICES restricted to one per proc).",
+    )
+    parser.add_argument(
+        "--fsdp",
+        default=False,
+        action="store_true",
+        help="Use FullyShardedDataParallel for distributed training.",
+    )
+    parser.add_argument(
+        "--fsdp_use_orig_params",
+        default=False,
+        action="store_true",
+        help="Passed into the FSDP constructor. Enables param_groups and gradient masking for weight_decay. Does not work with OPT.",
+    )
+    parser.add_argument(
+        "--fsdp_sharding_strategy", default="full", type=str, choices=["full", "hybrid"]
+    )
+
+    # wandb args
+    parser.add_argument("--report_to_wandb", default=False, action="store_true")
+    parser.add_argument(
+        "--wandb_project",
+        type=str,
+    )
+    parser.add_argument(
+        "--wandb_entity",
+        type=str,
+    )
+    parser.add_argument(
+        "--save_checkpoints_to_wandb",
+        default=False,
+        action="store_true",
+        help="save checkpoints to wandb",
+    )
+
+    args = parser.parse_args()
+
+    # Validate args
+    if args.laion_shards.startswith("s3"):
+        args.laion_shards = f"pipe:aws s3 cp {args.laion_shards} -"
+
+    if args.mmc4_shards.startswith("s3"):
+        args.mmc4_shards = f"pipe:aws s3 cp {args.mmc4_shards} -"
+
+    if args.save_checkpoints_to_wandb and not args.report_to_wandb:
+        raise ValueError("save_checkpoints_to_wandb requires report_to_wandb")
+
+    if args.fsdp and not args.fsdp_use_orig_params:
+        print(
+            "Warning: FSDP is running without fsdp_use_orig_params flag. "
+            + "This is not recommended because it means we will use uniform weight decay"
+            + " and train all embeddings, not just the newly added ones. "
+            + "Note: OPT models are not compatible with fsdp_use_orig_params flag."
+        )
+
+    if args.fsdp and args.fsdp_sharding_strategy == "hybrid":
+        print(
+            "Warning: As of torch=2.0.1, the FSDP logic for optim_state_dict() is broken for hybrid sharding."
+            + "To make this method work, we need to modify torch.distributed.fsdp._optim_utils.py"
+            + "Copy and paste the code from the _optim_utils.py in this repo into the torch file."
+            + "The main issue was the missing group kwarg on line 1596 in _all_gather_optim_state."
+        )
+
+    assert (args.train_num_samples_laion // args.batch_size_laion) == (
+        args.train_num_samples_mmc4 // args.batch_size_mmc4
+    ), "number of samples per epoch must be equal for mmc4 and laion"
+
+    # Set up distributed training
+    if args.offline:
+        os.environ["WANDB_MODE"] = "offline"
+        os.environ["TRANSFORMERS_OFFLINE"] = "1"
+    args.local_rank, args.rank, args.world_size = world_info_from_env()
+    device_id = init_distributed_device(args)
+    random_seed(args.seed)
+
+    # Initialize model
+    model, image_processor, tokenizer = create_model_and_transforms(
+        args.vision_encoder_path,
+        args.vision_encoder_pretrained,
+        args.lm_path,
+        args.tokenizer_path if args.tokenizer_path else args.lm_path,
+        cross_attn_every_n_layers=args.cross_attn_every_n_layers,
+        use_local_files=args.offline,
+        gradient_checkpointing=args.gradient_checkpointing,
+        freeze_lm_embeddings=args.freeze_lm_embeddings,
+    )
+    random_seed(args.seed, args.rank)
+
+    # Initialize logging
+    print(f"Start running training on rank {args.rank}.")
+    if args.rank == 0 and args.report_to_wandb:
+        wandb.init(
+            project=args.wandb_project,
+            entity=args.wandb_entity,
+            name=args.run_name,
+            config=vars(args),
+        )
+
+    # Load model checkpoint on CPU
+    if os.path.exists(f"{args.run_name}") and args.resume_from_checkpoint is None:
+        # if args do not specify a checkpoint to resume from, check if checkpoints exist for this run
+        # and automatically resume from the latest checkpoint
+        checkpoint_list = glob.glob(f"{args.run_name}/checkpoint_*.pt")
+        if len(checkpoint_list) == 0:
+            print(f"Found no checkpoints for run {args.run_name}.")
+        else:
+            args.resume_from_checkpoint = sorted(
+                checkpoint_list, key=lambda x: int(x.split("_")[-1].split(".")[0])
+            )[-1]
+            print(
+                f"Found checkpoint {args.resume_from_checkpoint} for run {args.run_name}."
+            )
+
+    resume_from_epoch = 0
+    if args.resume_from_checkpoint is not None:
+        if args.rank == 0:
+            print(f"Loading checkpoint from {args.resume_from_checkpoint}")
+        checkpoint = torch.load(args.resume_from_checkpoint, map_location="cpu")
+        msd = checkpoint["model_state_dict"]
+        msd = {k.replace("module.", ""): v for k, v in msd.items()}
+        resume_from_epoch = checkpoint["epoch"] + 1
+
+        # for fsdp, only one rank needs to load the state dict
+        if not args.fsdp or args.rank == 0:
+            model.load_state_dict(msd, False)
+
+    # Initialize FSDP / DDP, and ensure the model is on GPU
+    print(f"Initializing distributed training with {args.world_size} GPUs.")
+    if args.fsdp:
+        print(
+            f"Before FSDP parameter num: {sum(p.numel() for p in model.parameters())} on rank {args.rank}"
+        )
+
+        # init MixedPrecision
+        if args.precision != "fp32":
+            cast_dtype = get_mp_policy_dtype(args.precision)
+            mp_policy = MixedPrecision(
+                param_dtype=torch.float32,
+                reduce_dtype=cast_dtype,  # gradient communication
+                buffer_dtype=cast_dtype,
+            )
+        else:
+            mp_policy = None
+
+        # init process groups
+        if args.fsdp_sharding_strategy == "hybrid":
+            intra_node_group, inter_node_group = _init_intra_and_inter_node_groups(
+                _get_default_group()
+            )
+            args.my_group = intra_node_group  # for optimizer saving
+            process_group = (intra_node_group, inter_node_group)  # for FSDP init
+        else:
+            args.my_group = None  # for optimizer saving
+            process_group = None  # for FSDP init
+
+        # init FSDP
+        wrapper_kwargs = dict(
+            process_group=process_group,
+            cpu_offload=CPUOffload(offload_params=False),
+            device_id=device_id,
+            sync_module_states=True,  # broadcast loaded ckpt from rank 0 -> all ranks
+            sharding_strategy=ShardingStrategy.FULL_SHARD
+            if args.fsdp_sharding_strategy == "full"
+            else ShardingStrategy.HYBRID_SHARD,
+            use_orig_params=args.fsdp_use_orig_params,
+            mixed_precision=mp_policy,
+            forward_prefetch=True,
+            backward_prefetch=BackwardPrefetch.BACKWARD_PRE,
+            limit_all_gathers=True,
+        )
+        model.wrap_fsdp(wrapper_kwargs, device_id)
+        ddp_model = model
+
+        print(
+            f"After FSDP parameter num: {sum(p.numel() for p in model.parameters())} on rank {args.rank}"
+        )
+        print(
+            f"After FSDP {torch.cuda.memory_allocated()/1024**3:.3} GB on rank {args.rank}"
+        )
+
+    else:
+        model = model.to(device_id)
+        ddp_model = DDP(model, device_ids=[device_id])
+
+    # Initialize gradient checkpointing
+    if args.gradient_checkpointing:
+        non_reentrant_wrapper = functools.partial(
+            checkpoint_wrapper,
+            offload_to_cpu=True,
+            checkpoint_impl=CheckpointImpl.NO_REENTRANT,
+        )
+        apply_activation_checkpointing(
+            ddp_model,
+            checkpoint_wrapper_fn=non_reentrant_wrapper,
+            check_fn=lambda m: getattr(m, "_use_gradient_checkpointing", False)
+            and not isinstance(m, FSDP)
+            and not isinstance(m, CheckpointWrapper),
+        )
+
+    # Initialize optimizer
+    params_to_optimize = ddp_model.named_parameters()
+    params_to_optimize = list(
+        filter(
+            lambda x: x[1].requires_grad
+            and not getattr(x[1], "exclude_from_optimizer", False),
+            params_to_optimize,
+        )
+    )
+    if not args.fsdp or args.fsdp_use_orig_params:
+        # apply weight decay only to params in the xattn layers
+        def get_grouped_params(model):
+            params_with_wd, params_without_wd = [], []
+            for n, p in params_to_optimize:
+                if "gated_cross_attn" in n:
+                    params_with_wd.append(p)
+                else:
+                    params_without_wd.append(p)
+            return [
+                {"params": params_with_wd, "weight_decay": args.weight_decay},
+                {"params": params_without_wd, "weight_decay": 0.0},
+            ]
+
+        optimizer = torch.optim.AdamW(
+            get_grouped_params(params_to_optimize), lr=args.learning_rate
+        )
+    else:
+        # unclear if we should be using no weight decay or small weight decay for all parameters
+        optimizer = torch.optim.AdamW(
+            (p for _, p in params_to_optimize),
+            lr=args.learning_rate,
+            weight_decay=args.weight_decay,
+        )
+
+    # load optimizer checkpoint
+    if args.resume_from_checkpoint is not None:
+        osd = checkpoint["optimizer_state_dict"]
+        if args.fsdp:
+            osd = FSDP.optim_state_dict_to_load(osd, ddp_model, optimizer)
+        optimizer.load_state_dict(osd)
+
+    # Initialize data loaders
+    laion_dataset = get_data(args, image_processor, tokenizer, "image_text")
+    mmc4_dataset = get_data(args, image_processor, tokenizer, "mmc4")
+    total_training_steps = (
+        (args.train_num_samples_mmc4) // (args.batch_size_mmc4 * args.world_size)
+    ) * args.num_epochs
+
+    if args.rank == 0:
+        print(f"Total training steps: {total_training_steps}")
+
+    # Initialize lr scheduler
+    if args.lr_scheduler == "linear":
+        lr_scheduler = get_linear_schedule_with_warmup(
+            optimizer,
+            num_warmup_steps=args.warmup_steps,
+            num_training_steps=total_training_steps,
+        )
+    elif args.lr_scheduler == "cosine":
+        lr_scheduler = get_cosine_schedule_with_warmup(
+            optimizer,
+            num_warmup_steps=args.warmup_steps,
+            num_training_steps=total_training_steps,
+        )
+    else:
+        lr_scheduler = get_constant_schedule_with_warmup(
+            optimizer, num_warmup_steps=args.warmup_steps
+        )
+
+    # load lr scheduler checkpoint
+    if args.resume_from_checkpoint is not None:
+        lr_scheduler.load_state_dict(checkpoint["lr_scheduler_state_dict"])
+
+    # Start training!
+    ddp_model.train()
+
+    for epoch in range(resume_from_epoch, args.num_epochs):
+        laion_dataset.set_epoch(epoch)
+        laion_loader = laion_dataset.dataloader
+        mmc4_dataset.set_epoch(epoch)
+        mmc4_loader = mmc4_dataset.dataloader
+
+        train_one_epoch(
+            args=args,
+            model=ddp_model,
+            epoch=epoch,
+            tokenizer=tokenizer,
+            optimizer=optimizer,
+            lr_scheduler=lr_scheduler,
+            laion_loader=laion_loader,
+            mmc4_loader=mmc4_loader,
+            device_id=device_id,
+            wandb=wandb,
+        )
+        save_checkpoint(ddp_model, optimizer, lr_scheduler, epoch, args)
+
+    # save final checkpoint
+    save_checkpoint(ddp_model, optimizer, lr_scheduler, epoch, args)
+
+
+if __name__ == "__main__":
+    main()

open_flamingo/train/train_utils.py

@@ -0,0 +1,377 @@
+import time
+from contextlib import suppress
+import torch
+from tqdm import tqdm
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+from torch.distributed.fsdp import (
+    FullStateDictConfig,
+    StateDictType,
+)
+from torch.distributed.fsdp.api import FullOptimStateDictConfig
+import os
+import wandb
+from einops import rearrange
+
+
+def get_cast_dtype(precision: str):
+    cast_dtype = None
+    if precision == "bf16":
+        cast_dtype = torch.bfloat16
+    elif precision == "fp16":
+        cast_dtype = torch.float16
+    return cast_dtype
+
+
+def get_mp_policy_dtype(precision: str):
+    if "bfloat16" in precision or "bf16" in precision:
+        return torch.bfloat16
+    elif precision == "fp16":
+        return torch.float16
+    else:
+        return torch.float32
+
+
+def get_autocast(precision, cache_enabled=True):
+    if precision == "amp":
+        return torch.cuda.amp.autocast(cache_enabled=cache_enabled)
+    elif precision == "amp_bfloat16" or precision == "amp_bf16":
+        # amp_bfloat16 is more stable than amp float16 for clip training
+        return lambda: torch.cuda.amp.autocast(
+            dtype=torch.bfloat16, cache_enabled=cache_enabled
+        )
+    else:
+        return suppress
+
+
+def train_one_epoch(
+    args,
+    model,
+    epoch,
+    laion_loader,
+    mmc4_loader,
+    tokenizer,
+    optimizer,
+    lr_scheduler,
+    device_id,
+    wandb,
+):
+    # setup loaders
+    num_batches_per_epoch_laion = laion_loader.num_batches
+    num_batches_per_epoch_mmc4 = mmc4_loader.num_batches
+    assert (
+        num_batches_per_epoch_laion == num_batches_per_epoch_mmc4
+    ), "Number of batches in laion and mmc4 datasets must be the same"
+    num_batches_per_epoch = num_batches_per_epoch_mmc4
+    total_training_steps = num_batches_per_epoch * args.num_epochs
+
+    autocast = get_autocast(
+        args.precision, cache_enabled=(not args.fsdp)
+    )  # if fsdp, disable cache to save memory
+    cast_dtype = get_cast_dtype(args.precision)
+
+    # setup model
+    media_token_id = tokenizer("<image>", add_special_tokens=False)["input_ids"][-1]
+    endofchunk_token_id = tokenizer("<|endofchunk|>", add_special_tokens=False)[
+        "input_ids"
+    ][-1]
+    model.train()
+
+    # setup logging
+    step_time_m = AverageMeter()
+    data_time_m = AverageMeter()
+    end = time.time()
+
+    # loop through dataloader
+    for num_steps, (batch_laion, batch_mmc4) in tqdm(
+        enumerate(zip(laion_loader, mmc4_loader)),
+        disable=args.rank != 0,
+        total=total_training_steps,
+        initial=(epoch * num_batches_per_epoch),
+    ):
+        data_time_m.update(time.time() - end)
+        global_step = num_steps + epoch * num_batches_per_epoch
+
+        #### LAION FORWARD PASS ####
+        images = batch_laion[0].to(device_id, dtype=cast_dtype, non_blocking=True)
+        images = rearrange(images, "(b t f) c h w -> b t f c h w", t=1, f=1)
+        input_ids = batch_laion[1][0].to(device_id, dtype=cast_dtype, non_blocking=True)
+        attention_mask = batch_laion[1][1].to(
+            device_id, dtype=cast_dtype, non_blocking=True
+        )
+
+        # set up labels; language model is expected to handle shifting
+        labels = input_ids.clone()
+        labels[labels == tokenizer.pad_token_id] = -100
+        labels[labels == tokenizer.eos_token] = -100
+        labels[labels == media_token_id] = -100
+        labels = labels.to(device_id)
+
+        # gradient accumulation w/ fsdp cpu offloading requires a no_sync context manager
+        with autocast():
+            loss_laion = model(
+                vision_x=images,
+                lang_x=input_ids,
+                attention_mask=attention_mask,
+                labels=labels,
+            )[0]
+
+        divided_loss_laion = loss_laion / args.gradient_accumulation_steps
+        (divided_loss_laion * args.loss_multiplier_laion).backward()
+
+        #### MMC4 FORWARD PASS ####
+        images = batch_mmc4[0].to(device_id, dtype=cast_dtype, non_blocking=True)
+        images = rearrange(images, "b (t f) c h w -> b t f c h w", f=1)
+        input_ids = torch.stack([x[0] for x in batch_mmc4[1]]).squeeze(1)
+        attention_mask = torch.stack([x[1] for x in batch_mmc4[1]]).squeeze(1)
+
+        # set up labels; language model is expected to handle shifting
+        labels = input_ids.clone()
+        labels[labels == tokenizer.pad_token_id] = -100
+        labels[labels == tokenizer.eos_token] = -100
+        for i in range(labels.shape[0]):
+            # remove loss for any token before the first <image> token
+            label_idx = 0
+            while (
+                label_idx < labels.shape[1] and labels[i][label_idx] != media_token_id
+            ):
+                labels[i][label_idx] = -100
+                label_idx += 1
+
+            # get index of all endofchunk tokens in the sequence
+            endofchunk_idxs = torch.where(labels[i] == endofchunk_token_id)[0]
+            for endofchunk_idx in endofchunk_idxs:
+                token_idx = endofchunk_idx + 1
+                while (
+                    token_idx < labels.shape[1]
+                    and labels[i][token_idx] != media_token_id
+                ):
+                    labels[i][token_idx] = -100
+                    token_idx += 1
+
+        labels[labels == media_token_id] = -100
+        labels = labels.to(device_id)
+
+        # gradient accumulation w/ fsdp cpu offloading requires a no_sync context manager
+        with autocast():
+            loss_mmc4 = model(
+                vision_x=images,
+                lang_x=input_ids.to(device_id),
+                attention_mask=attention_mask.to(device_id),
+                labels=labels,
+            )[0]
+
+            # if loss is nan, skip this batch
+            # this hack of skipping the batch is not FSDP-compatible
+            if torch.isnan(loss_mmc4):
+                print("loss is nan, skipping this batch")
+                print("input_ids: ", tokenizer.batch_decode(input_ids))
+                print("labels: ", labels)
+                print("images: ", images)
+                optimizer.zero_grad(set_to_none=True)
+                continue
+
+        divided_loss_mmc4 = loss_mmc4 / args.gradient_accumulation_steps
+        (divided_loss_mmc4 * args.loss_multiplier_mmc4).backward()
+
+        if (not args.freeze_lm_embeddings) and (
+            not args.fsdp or args.fsdp_use_orig_params
+        ):
+            # Mask gradients for input embeddings s.t. we only update the added tokens <image> and <|endofchunk|>
+            if args.fsdp:
+                embed_grad = model.lang_encoder.get_input_embeddings().weight.grad
+            else:
+                embed_grad = (
+                    model.module.lang_encoder.get_input_embeddings().weight.grad
+                )
+            zero_mask = torch.zeros_like(embed_grad)
+            zero_mask[media_token_id] = torch.ones_like(zero_mask[media_token_id])
+            zero_mask[endofchunk_token_id] = torch.ones_like(
+                zero_mask[endofchunk_token_id]
+            )
+            if args.fsdp:
+                model.lang_encoder.get_input_embeddings().weight.grad = (
+                    embed_grad * zero_mask
+                )
+            else:
+                model.module.lang_encoder.get_input_embeddings().weight.grad = (
+                    embed_grad * zero_mask
+                )
+
+        # clip gradient norm
+        if args.fsdp:
+            """
+            The way we clip gradients with FSDP is different than the non-FSDP case,
+            because during FSDP, gradient norms are computed over certain submodules,
+            rather than the entire model.
+            At least for OPT-125M, this didn't seem to make a difference in performance.
+            """
+            model.clip_grad_norm_(1.0)
+        else:
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+
+        # step optimizer and log
+        if (((num_steps + 1) % args.gradient_accumulation_steps) == 0) or (
+            num_steps == num_batches_per_epoch - 1
+        ):
+            optimizer.step()
+            lr_scheduler.step()
+            optimizer.zero_grad(set_to_none=True)
+
+            # step time and reset end outside of rank 0
+            step_time_m.update(time.time() - end)
+            end = time.time()
+
+            # rank 0 logging
+            if args.rank == 0 and args.report_to_wandb:
+                laion_samples_per_second = (
+                    args.gradient_accumulation_steps
+                    * args.batch_size_laion
+                    * args.world_size
+                    / step_time_m.val
+                )
+                laion_samples_per_second_per_gpu = (
+                    args.gradient_accumulation_steps
+                    * args.batch_size_laion
+                    / step_time_m.val
+                )
+                c4_samples_per_second = (
+                    args.gradient_accumulation_steps
+                    * args.batch_size_mmc4
+                    * args.world_size
+                    / step_time_m.val
+                )
+                c4_samples_per_second_per_gpu = (
+                    args.gradient_accumulation_steps
+                    * args.batch_size_mmc4
+                    / step_time_m.val
+                )
+                wandb.log(
+                    {
+                        "data_time": data_time_m.avg,
+                        "step_time": step_time_m.avg,
+                        "laion_samples_per_second": laion_samples_per_second,
+                        "laion_samples_per_second_per_gpu": laion_samples_per_second_per_gpu,
+                        "c4_samples_per_second": c4_samples_per_second,
+                        "c4_samples_per_second_per_gpu": c4_samples_per_second_per_gpu,
+                        "lr": optimizer.param_groups[0]["lr"],
+                    },
+                    commit=False,
+                )
+                step_time_m.reset()
+                data_time_m.reset()
+
+                wandb.log(
+                    {
+                        "loss_laion": loss_laion.item(),
+                        "global_step": global_step,
+                    },
+                    commit=False,
+                )
+                wandb.log(
+                    {"loss_mmc4": loss_mmc4.item(), "global_step": global_step},
+                    commit=True,
+                )
+
+        # Log loss to console
+        if ((num_steps + 1) % args.logging_steps == 0) and args.rank == 0:
+            print(
+                f"Step {num_steps+1}/{num_batches_per_epoch} of epoch {epoch+1}/{args.num_epochs} complete. Loss LAION: {loss_laion.item():.3f} // Loss MMC4: {loss_mmc4.item():.3f}"
+            )
+
+
+class AverageMeter(object):
+    """Computes and stores the average and current value"""
+
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        self.val = 0
+        self.avg = 0
+        self.sum = 0
+        self.count = 0
+
+    def update(self, val, n=1):
+        self.val = val
+        self.sum += val * n
+        self.count += n
+        self.avg = self.sum / self.count
+
+
+def filter_state_dict_to_trainable(model, state_dict):
+    """
+    Remove non-trainable parameters from model state dict.
+    Exception: Embeddings will not be removed, even if frozen.
+    This is because we need the new <image> <|endofchunk|> tokens to
+    be consistent across initializations.
+    """
+    for (
+        name,
+        p,
+    ) in model.named_parameters():  # won't work for fsdp + use_orig_params=False
+        if "fsdp" in name:
+            continue
+        if "embed" in name or isinstance(p, torch.nn.Embedding):
+            continue
+        if not p.requires_grad:
+            name = name.replace("._checkpoint_wrapped_module", "")
+            if name in state_dict:
+                del state_dict[name]
+            else:
+                print(f"WARNING: filtering but {name} not in state_dict")
+
+    # also remove the keys in state_dict generated from
+    # lang_encoder.old_decoder_blocks and lang_encoder.gated_cross_attn_layers
+    # because these are already saved in lang_encoder.model...
+    to_delete = [
+        n
+        for n in state_dict.keys()
+        if ("lang_encoder.old_decoder_blocks" in n)
+        or ("lang_encoder.gated_cross_attn_layers" in n)
+        or ("vision_encoder" in n)
+    ]
+    for name in to_delete:
+        del state_dict[name]
+    return state_dict
+
+
+def save_checkpoint(model, optimizer, lr_scheduler, epoch, args):
+    """
+    Save training checkpoint with model, optimizer, and lr_scheduler state.
+    """
+    if args.fsdp:
+        FSDP.set_state_dict_type(
+            model,
+            StateDictType.FULL_STATE_DICT,
+            FullStateDictConfig(rank0_only=True, offload_to_cpu=True),
+            FullOptimStateDictConfig(rank0_only=True),
+        )
+        model_state = model.state_dict()
+        optim_state = FSDP.optim_state_dict(model, optimizer, group=args.my_group)
+
+    else:
+        model_state = model.state_dict()
+        optim_state = optimizer.state_dict()
+
+    if args.rank == 0:
+        if not (args.fsdp and not args.fsdp_use_orig_params):
+            model_state = filter_state_dict_to_trainable(model, model_state)
+
+        if not os.path.exists(args.run_name):
+            os.makedirs(args.run_name)
+
+        checkpoint_dict = {
+            "epoch": epoch,
+            "model_state_dict": model_state,
+            "optimizer_state_dict": optim_state,
+            "lr_scheduler_state_dict": lr_scheduler.state_dict(),
+        }
+
+        print(f"Saving checkpoint to {args.run_name}/checkpoint_{epoch}.pt")
+        torch.save(checkpoint_dict, f"{args.run_name}/checkpoint_{epoch}.pt")
+        if args.report_to_wandb and args.save_checkpoints_to_wandb:
+            wandb.save(f"{args.run_name}/checkpoint_{epoch}.pt")
+
+        if args.delete_previous_checkpoint:
+            if epoch > 0:
+                os.remove(f"{args.run_name}/checkpoint_{epoch-1}.pt")

requirements.txt

@@ -4,14 +4,9 @@ torchvision==0.16.0
 transformers==4.35.0
 huggingface_hub>=0.16.4,<1.0
 einops==0.7.0
-einops-exts==0.0.4
 numpy==1.24.3
 Pillow==10.1.0
 matplotlib==3.8.0
 open_clip_torch==2.23.0
 accelerate==0.24.1
-safetensors==0.4.0
-zarr==2.16.1
-numcodecs==0.12.1
-hydra-core==1.3.2
-omegaconf==2.3.0
+safetensors==0.4.0