Multimodal Training Data › Code Implementation: Processing Multimodal Datasets

import subprocess
import sys

subprocess.run(
    [
        "uv",
        "pip",
        "install",
        "--python",
        sys.executable,
        "datasets",
        "transformers",
        "torch",
        "requests",
        "Pillow",
        "matplotlib",
        "numpy",
        "-q",
    ],
    check=True,
)

# Load a small sample of Conceptual Captions 12M
# In practice, use streaming=True for the full dataset
try:
    dataset = load_dataset(
        "conceptual_captions", "unlabeled", split="train", streaming=True
    )
    # Take first 100 examples for demonstration
    sample_data = []
    for i, example in enumerate(dataset):
        if i >= 100:
            break
        sample_data.append(example)
except Exception:
    # Create mock data for demonstration when dataset is unavailable
    captions = [
        "A red apple on a wooden table",
        "A golden retriever playing in a park",
        "A vintage bicycle leaning against a brick wall",
        "A steaming cup of coffee on a rainy morning",
        "Colorful hot air balloons floating over mountains",
        "An ancient stone archway in a medieval town",
        "A modern glass skyscraper which shows the sunset",
        "Fresh strawberries in a ceramic bowl",
    ]
    sample_data = [
        {
            "image_url": f"http://example.com/{i}.jpg",
            "caption": captions[i % len(captions)],
        }
        for i in range(20)
    ]

import hashlib
from io import BytesIO

import requests
from PIL import Image


def load_image_from_url(url, timeout=5):
    """Download and convert image URL to PIL Image."""
    try:
        response = requests.get(url, timeout=timeout)
        response.raise_for_status()
        image = Image.open(BytesIO(response.content)).convert("RGB")
        return image
    except Exception:
        # Generate a synthetic image as fallback when URL fails
        fallback_seed = int.from_bytes(
            hashlib.sha256(url.encode("utf-8")).digest()[:4], "little"
        )
        fallback_rng = np.random.default_rng(fallback_seed)
        img_array = fallback_rng.integers(0, 255, (224, 224, 3), dtype=np.uint8)
        return Image.fromarray(img_array)


# Attempt to load images for our sample
loaded_examples = []
for ex in sample_data[:20]:  # Try first 20 to account for broken links
    img = load_image_from_url(ex["image_url"])
    if img is not None:
        loaded_examples.append(
            {"image": img, "caption": ex["caption"], "url": ex["image_url"]}
        )

import torch
from transformers import CLIPModel, CLIPProcessor

# Load a small CLIP model for similarity scoring
model_name = "openai/clip-vit-base-patch32"
processor = CLIPProcessor.from_pretrained(model_name)
model = CLIPModel.from_pretrained(model_name)


def compute_clip_similarity(image, text):
    """Compute CLIP cosine similarity between image and text."""
    inputs = processor(
        text=[text], images=image, return_tensors="pt", padding=True
    )

    with torch.no_grad():
        outputs = model(**inputs)
        image_embeds = outputs.image_embeds
        text_embeds = outputs.text_embeds

        # Normalize and compute cosine similarity
        image_embeds = image_embeds / image_embeds.norm(dim=-1, keepdim=True)
        text_embeds = text_embeds / text_embeds.norm(dim=-1, keepdim=True)
        similarity = (image_embeds @ text_embeds.T).item()

    return similarity

# Compute similarities for the six examples displayed below
similarities = []
for i, ex in enumerate(loaded_examples[:6]):
    sim = compute_clip_similarity(ex["image"], ex["caption"])
    similarities.append(sim)

def create_interleaved_sequence(examples, max_images=3):
    """
    Simulate an interleaved document by combining multiple image-text pairs
    into a sequence with special tokens.
    """
    selected = examples[:max_images]

    sequence = []
    for i, ex in enumerate(selected):
        # Add text context
        sequence.append(f"<text>Figure {i + 1} shows: {ex['caption']}</text>")
        # Add image placeholder
        sequence.append("<image>")

    return {
        "images": [ex["image"] for ex in selected],
        "text_sequence": " ".join(sequence),
        "num_images": len(selected),
    }


interleaved = create_interleaved_sequence(loaded_examples, max_images=3)

def collate_multimodal_batch(examples, proc, max_length=128):
    """
    Collate function for multimodal training batches.
    Handles variable image counts and text lengths.
    """
    images = []
    texts = []

    for ex in examples:
        # Each example may have multiple images
        imgs = ex.get("images", [ex.get("image")])
        images.extend(imgs)

        # Create text with image tokens
        text = ex.get("text_sequence", ex.get("caption"))
        # Replace <image> with special tokens or process separately
        texts.append(text)

    # Process all images together
    if images:
        pixel_values = proc(images=images, return_tensors="pt")["pixel_values"]
    else:
        pixel_values = None

    # Process text
    text_inputs = proc(
        text=texts,
        return_tensors="pt",
        padding=True,
        truncation=True,
        max_length=max_length,
    )

    return {
        "pixel_values": pixel_values,
        "input_ids": text_inputs["input_ids"],
        "attention_mask": text_inputs["attention_mask"],
    }

def analyze_filtering_thresholds(
    examples, thresholds=[0.1, 0.15, 0.2, 0.25, 0.3, 0.35]
):
    """
    Compute retention rates at different CLIP similarity thresholds.
    """
    # Compute all similarities first
    all_sims = []
    for ex in examples:
        sim = compute_clip_similarity(ex["image"], ex["caption"])
        all_sims.append(sim)

    all_sims = np.array(all_sims)
    results = {}
    for thresh in thresholds:
        retained = np.sum(all_sims >= thresh)
        results[thresh] = {
            "retained": int(retained),
            "total": len(all_sims),
            "fraction": retained / len(all_sims),
            "mean_sim_retained": float(all_sims[all_sims >= thresh].mean())
            if retained > 0
            else 0.0,
        }
    return results, all_sims


# Analyze our sample (may be synthetic images with low CLIP scores)
filter_results, all_similarities = analyze_filtering_thresholds(
    loaded_examples[:10]
)