CLIP and SigLIP put images and text in one space

import importlib.util
import subprocess
import sys

TFS = "git+https://github.com/xmarva/transformers-from-scratch@main"
if importlib.util.find_spec("tfs") is None:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", TFS])

from tfs import env
from tfs.viz import table, token_chips, use_course_style

env.quiet()  # no progress bars or environment warnings in the outputs
SEED = 42
env.seed_everything(SEED)
COLORS = use_course_style()
WORK_DIR = env.work_dir()

import io
import math
import tarfile
import textwrap
import time
import urllib.request
from collections import Counter
from concurrent.futures import ThreadPoolExecutor

import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap
from matplotlib.offsetbox import AnnotationBbox, OffsetImage
from matplotlib import font_manager
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
from IPython.display import display
from PIL import Image, ImageDraw, ImageFont
from transformers import AutoImageProcessor, AutoModel, AutoProcessor, AutoTokenizer
from transformers import CLIPImageProcessor, CLIPModel, CLIPTokenizer

SMOKE = env.smoke()  # CI runs the lesson on a small slice; you run it in full
device = "cuda" if torch.cuda.is_available() else "cpu"
if device == "cuda":
    print("GPU:", torch.cuda.get_device_name(0))

# images are teal, captions amber, the pairs that belong together rose, as in the figures
heat = LinearSegmentedColormap.from_list("heat", ["#FFFFFF", COLORS["dec"]])
diverging = LinearSegmentedColormap.from_list(
    "diverging", [COLORS["enc"], "#FFFFFF", COLORS["dec"]]
)

HUB = "https://huggingface.co/datasets/clip-benchmark/"


def examples(url, limit=None):
    """Yield (name, {extension: bytes}) for each example of a shard, as the bytes arrive."""
    count = 0
    with urllib.request.urlopen(url, timeout=120) as response:
        archive = tarfile.open(fileobj=response, mode="r|")  # "|": a stream, read once
        name, files = None, {}
        for member in archive:
            stem, extension = member.name.rsplit(".", 1)
            if stem != name and files:
                yield name, files
                count += 1
                if limit and count >= limit:
                    return
                files = {}
            name = stem
            files[extension] = archive.extractfile(member).read()
        if files:
            yield name, files


def thumbnail(data, side=256):
    photo = Image.open(io.BytesIO(data)).convert("RGB")
    scale = side / min(photo.size)
    if scale < 1:
        new_size = (round(photo.width * scale), round(photo.height * scale))
        photo = photo.resize(new_size, Image.BICUBIC)
    buffer = io.BytesIO()
    photo.save(buffer, "JPEG", quality=90)
    return buffer.getvalue()


def read_shards(urls, limit, label):
    """Stream the shards; make thumbnails in a pool of threads while the bytes arrive."""
    pending = []
    with ThreadPoolExecutor(8) as pool:
        for url in urls:
            for _, files in examples(url, limit):
                pending.append((pool.submit(thumbnail, files["jpg"]), label(files)))
        return [(job.result(), value) for job, value in pending]


def load(thumb):
    return Image.open(io.BytesIO(thumb)).convert("RGB")

def text_file(path):
    return urllib.request.urlopen(HUB + path).read().decode().strip().split("\n")


COUNTRIES = text_file("wds_country211/resolve/main/classnames.txt")
PAPER_TEMPLATES = text_file(
    "wds_country211/resolve/main/zeroshot_classification_templates.txt"
)
start = time.time()
geo = read_shards(
    [HUB + f"wds_country211/resolve/main/test/{k}.tar" for k in range(4)],
    limit=30 if SMOKE else None,  # the shards are sorted by country
    label=lambda files: int(files["cls"]),
)
geo_thumbs = [thumb for thumb, _ in geo]
geo_truth = torch.tensor([country for _, country in geo])
print(len(COUNTRIES), "countries:", ", ".join(COUNTRIES[:6]), "...")
print(f"{len(geo_thumbs)} photos in {time.time() - start:.0f} s")
print("the paper's sentences:", PAPER_TEMPLATES)

start = time.time()
coco = read_shards(
    [HUB + f"wds_mscoco_captions/resolve/main/test/{k}.tar" for k in range(5)],
    limit=64 if SMOKE else None,
    label=lambda files: [line.strip() for line in files["txt"].decode().split("\n")],
)
coco_thumbs = [thumb for thumb, _ in coco]
coco_captions = [[c for c in captions if c] for _, captions in coco]
captions = [c for photo_captions in coco_captions for c in photo_captions]
owner = torch.tensor(
    [i for i, cs in enumerate(coco_captions) for _ in cs]
)  # caption -> photo
first = torch.tensor([0] + np.cumsum([len(cs) for cs in coco_captions])[:-1].tolist())
BATCH = [26, 51, 27, 37, 50, 10, 23, 1]  # the eight photos the lesson comes back to
print(
    f"{len(coco_thumbs)} photos, {len(captions)} captions, in {time.time() - start:.0f} s"
)
size = sum(map(len, geo_thumbs + coco_thumbs)) / 1e6
print(f"all the thumbnails take {size:.0f} MB of memory")

fig, axes = plt.subplots(2, 4, figsize=(12, 7))
SHOW = [
    (0, 0),
    (53, 24),
    (106, 18),
    (159, 7),
]  # (country, its k-th photo): places, no faces
for ax, (country, k) in zip(axes[0], SHOW):
    i = (geo_truth == country).nonzero()[k].item()
    ax.imshow(load(geo_thumbs[i]))
    ax.set_title(COUNTRIES[country], fontsize=10)
    ax.axis("off")
for ax, i in zip(axes[1], BATCH[:4]):
    ax.imshow(load(coco_thumbs[i]))
    ax.set_title(textwrap.fill(coco_captions[i][0], 34), fontsize=8, loc="left")
    ax.axis("off")
plt.tight_layout()
plt.show()

CLIP = "openai/clip-vit-base-patch32"
clip_hf = CLIPModel.from_pretrained(CLIP).eval()
tokenizer = CLIPTokenizer.from_pretrained(CLIP)
processor = CLIPImageProcessor.from_pretrained(CLIP)
weights = clip_hf.state_dict()
vision_config = clip_hf.config.vision_config
text_config = clip_hf.config.text_config
print(
    "image tower:",
    vision_config.num_hidden_layers,
    "blocks of",
    vision_config.hidden_size,
)
print(
    "text tower:", text_config.num_hidden_layers, "blocks of", text_config.hidden_size
)
print("common space:", clip_hf.config.projection_dim, "numbers")

EXAMPLE = BATCH[0]
photo = load(coco_thumbs[EXAMPLE])
caption = coco_captions[EXAMPLE][0]
print("the caption:", caption)

SIZE = vision_config.image_size
MEAN = torch.tensor(processor.image_mean).reshape(3, 1, 1)
STD = torch.tensor(processor.image_std).reshape(3, 1, 1)
print("mean:", processor.image_mean)
print("std:", processor.image_std)


def clip_pixels(photo, size=SIZE):
    width, height = photo.size
    scale = size / min(width, height)
    new_size = (max(size, int(width * scale)), max(size, int(height * scale)))
    photo = photo.resize(new_size, Image.BICUBIC)
    left = (photo.width - size) // 2
    top = (photo.height - size) // 2
    photo = photo.crop((left, top, left + size, top + size))
    pixels = torch.tensor(np.asarray(photo), dtype=torch.float32).permute(2, 0, 1) / 255
    return (pixels - MEAN) / STD


x = clip_pixels(photo)
reference = processor(images=photo, return_tensors="pt")["pixel_values"][0]
print("our tensor:", tuple(x.shape))
levels = ((x - reference).abs() * STD * 255).max().item()  # back to 0..255 brightness
print(
    "largest difference from Hugging Face's processor:",
    round(levels, 2),
    "of 255 levels",
)

PATCH = vision_config.patch_size
grid = SIZE // PATCH
patches = (
    x.reshape(3, grid, PATCH, grid, PATCH)
    .permute(1, 3, 0, 2, 4)
    .reshape(grid * grid, -1)
)
print("patches:", tuple(patches.shape))

kernel = weights["vision_model.embeddings.patch_embedding.weight"]  # (768, 3, 32, 32)
by_matrix = patches @ kernel.reshape(kernel.shape[0], -1).T
by_conv = F.conv2d(x[None], kernel, stride=PATCH)[0].flatten(1).T
print("projected patches:", tuple(by_matrix.shape))
print(
    "matrix against convolution, largest difference:",
    (by_matrix - by_conv).abs().max().item(),
)

class_vector = weights["vision_model.embeddings.class_embedding"]
positions = weights["vision_model.embeddings.position_embedding.weight"]
sequence = torch.cat([class_vector[None], by_matrix]) + positions
print("class token + 49 patches, with positions:", tuple(sequence.shape))
sequence = F.layer_norm(
    sequence,
    (sequence.shape[-1],),
    weights["vision_model.pre_layrnorm.weight"],
    weights["vision_model.pre_layrnorm.bias"],
)

with torch.no_grad():
    theirs = clip_hf.vision_model.embeddings(x[None])
    theirs = clip_hf.vision_model.pre_layrnorm(theirs)
    tower = clip_hf.vision_model(pixel_values=x[None])
    image_vector = clip_hf.visual_projection(tower.pooler_output)[0]
print(
    "our sequence against the model's, largest difference:",
    (sequence - theirs[0]).abs().max().item(),
)
print("after the blocks:", tuple(tower.last_hidden_state.shape))
print("the class token's output, normalised:", tuple(tower.pooler_output.shape))
print("the image's vector:", tuple(image_vector.shape))

encoded = tokenizer([caption], truncation=True, return_tensors="pt")
ids = encoded["input_ids"][0]
pieces = tokenizer.convert_ids_to_tokens(ids)
token_chips(
    [pieces],
    ids=[ids.tolist()],
    strong=lambda row, col, token: token == "<|endoftext|>",
    mark_spaces=False,
    max_width=700,
)
print(len(ids), "tokens, the start and end markers included")

token_table = weights["text_model.embeddings.token_embedding.weight"]
text_positions = weights["text_model.embeddings.position_embedding.weight"]
n = len(ids)
text_sequence = token_table[ids] + text_positions[:n]
print("token table:", tuple(token_table.shape))
print("the caption as a sequence:", tuple(text_sequence.shape))

causal = torch.full((n, n), float("-inf")).triu(1)
fig, ax = plt.subplots(figsize=(4.2, 4.2))
ax.imshow(
    torch.isfinite(causal),
    cmap=LinearSegmentedColormap.from_list("m", ["#FFFFFF", COLORS["att"]]),
)
ax.set_xticks(range(n), pieces, rotation=90, fontsize=8)
ax.set_yticks(range(n), pieces, fontsize=8)
ax.set_xlabel("may look at")
ax.set_title("which tokens each token may see")
ax.grid(False)
plt.show()

with torch.no_grad():
    theirs = clip_hf.text_model.embeddings(input_ids=ids[None])
    tower = clip_hf.text_model(input_ids=ids[None])
    end = ids.argmax()
    at_end = tower.last_hidden_state[0, end]
    text_vector = clip_hf.text_projection(tower.pooler_output)[0]
print(
    "our sequence against the model's, largest difference:",
    (text_sequence - theirs[0]).abs().max().item(),
)
print("the end token sits at position", end.item(), "of", n)
print(
    "the model's summary is the output at that position:",
    torch.equal(at_end, tower.pooler_output[0]),
)
print("the caption's vector:", tuple(text_vector.shape))

others = [coco_captions[i][0] for i in BATCH]
with torch.no_grad():
    batch = tokenizer(others, padding=True, truncation=True, return_tensors="pt")
    other_vectors = clip_hf.text_projection(clip_hf.text_model(**batch).pooler_output)
image_unit = F.normalize(image_vector, dim=-1)
text_units = F.normalize(other_vectors, dim=-1)
cosines = text_units @ image_unit
rows = []
for text, value in zip(others, cosines.tolist()):
    rows.append({"caption": text[:70], "cosine with the photo": value})
table(pd.DataFrame(rows), digits=3)

def quick_gelu(x):
    return x * torch.sigmoid(1.702 * x)


class Attention(nn.Module):
    def __init__(self, d, heads):
        super().__init__()
        self.heads = heads
        self.head_dim = d // heads
        self.query = nn.Linear(d, d)
        self.key = nn.Linear(d, d)
        self.value = nn.Linear(d, d)
        self.out = nn.Linear(d, d)

    def forward(self, x, mask=None):
        b, n, d = x.shape
        q = self.query(x).view(b, n, self.heads, self.head_dim).transpose(1, 2)
        k = self.key(x).view(b, n, self.heads, self.head_dim).transpose(1, 2)
        v = self.value(x).view(b, n, self.heads, self.head_dim).transpose(1, 2)
        out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
        return self.out(out.transpose(1, 2).reshape(b, n, d))


class Block(nn.Module):
    def __init__(self, d, heads, hidden):
        super().__init__()
        self.norm1 = nn.LayerNorm(d)
        self.attn = Attention(d, heads)
        self.norm2 = nn.LayerNorm(d)
        self.fc1 = nn.Linear(d, hidden)
        self.fc2 = nn.Linear(hidden, d)

    def forward(self, x, mask=None):
        x = x + self.attn(self.norm1(x), mask)
        return x + self.fc2(quick_gelu(self.fc1(self.norm2(x))))

class ImageTower(nn.Module):
    def __init__(self, c, out):
        super().__init__()
        d, p = c.hidden_size, c.patch_size
        self.patch_embed = nn.Conv2d(3, d, kernel_size=p, stride=p, bias=False)
        self.cls = nn.Parameter(torch.zeros(d))
        self.pos = nn.Parameter(torch.zeros((c.image_size // p) ** 2 + 1, d))
        self.norm_pre = nn.LayerNorm(d)
        blocks = [
            Block(d, c.num_attention_heads, c.intermediate_size)
            for _ in range(c.num_hidden_layers)
        ]
        self.blocks = nn.ModuleList(blocks)
        self.norm_post = nn.LayerNorm(d)
        self.proj = nn.Linear(d, out, bias=False)

    def forward(self, pixels):
        x = self.patch_embed(pixels).flatten(2).transpose(1, 2)
        x = torch.cat([self.cls.expand(x.shape[0], 1, -1), x], dim=1) + self.pos
        x = self.norm_pre(x)
        for block in self.blocks:
            x = block(x)
        return self.proj(self.norm_post(x[:, 0]))


class TextTower(nn.Module):
    def __init__(self, c, out):
        super().__init__()
        d = c.hidden_size
        self.tok = nn.Embedding(c.vocab_size, d)
        self.pos = nn.Parameter(torch.zeros(c.max_position_embeddings, d))
        blocks = [
            Block(d, c.num_attention_heads, c.intermediate_size)
            for _ in range(c.num_hidden_layers)
        ]
        self.blocks = nn.ModuleList(blocks)
        self.norm = nn.LayerNorm(d)
        self.proj = nn.Linear(d, out, bias=False)

    def forward(self, ids, attention_mask):
        n = ids.shape[1]
        x = self.tok(ids) + self.pos[:n]
        causal = torch.full((n, n), float("-inf"), device=ids.device).triu(1)
        padding = torch.zeros(attention_mask.shape, device=ids.device)
        padding = padding.masked_fill(attention_mask == 0, float("-inf"))
        mask = causal[None, None] + padding[:, None, None, :]
        for block in self.blocks:
            x = block(x, mask)
        x = self.norm(x)
        end = ids.argmax(dim=-1)  # the end token has the largest id
        return self.proj(x[torch.arange(len(ids)), end])


class CLIP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.image = ImageTower(config.vision_config, config.projection_dim)
        self.text = TextTower(config.text_config, config.projection_dim)
        self.logit_scale = nn.Parameter(torch.zeros(()))

def rename(state):
    top = {
        "vision_model.embeddings.patch_embedding": "image.patch_embed",
        "vision_model.embeddings.class_embedding": "image.cls",
        "vision_model.embeddings.position_embedding.weight": "image.pos",
        "vision_model.pre_layrnorm": "image.norm_pre",
        "vision_model.post_layernorm": "image.norm_post",
        "visual_projection": "image.proj",
        "text_model.embeddings.token_embedding": "text.tok",
        "text_model.embeddings.position_embedding.weight": "text.pos",
        "text_model.final_layer_norm": "text.norm",
        "text_projection": "text.proj",
        "logit_scale": "logit_scale",
    }
    inside = {
        "self_attn.q_proj": "attn.query",
        "self_attn.k_proj": "attn.key",
        "self_attn.v_proj": "attn.value",
        "self_attn.out_proj": "attn.out",
        "layer_norm1": "norm1",
        "layer_norm2": "norm2",
        "mlp.fc1": "fc1",
        "mlp.fc2": "fc2",
    }
    out = {}
    for key, value in state.items():
        if "position_ids" in key:
            continue  # a buffer of 0, 1, 2, ..., not a weight
        for tower, ours in [
            ("vision_model.encoder.layers.", "image"),
            ("text_model.encoder.layers.", "text"),
        ]:
            if key.startswith(tower):
                number, rest = key.removeprefix(tower).split(".", 1)
                for theirs, mine in inside.items():
                    rest = rest.replace(theirs, mine)
                key = f"{ours}.blocks.{number}.{rest}"
        for theirs, mine in top.items():
            if key.startswith(theirs):
                key = mine + key[len(theirs) :]
        out[key] = value
    return out


clip = CLIP(clip_hf.config)
clip.load_state_dict(rename(weights), strict=True)
clip = clip.to(device).eval()
clip_hf = clip_hf.to(device)

check = [load(coco_thumbs[i]) for i in BATCH]
with torch.no_grad():
    pixels = torch.stack([clip_pixels(photo) for photo in check]).to(device)
    batch = tokenizer(others, padding=True, truncation=True, return_tensors="pt").to(
        device
    )
    ours_image = clip.image(pixels)
    ours_text = clip.text(batch["input_ids"], batch["attention_mask"])
    theirs_image = clip_hf.visual_projection(
        clip_hf.vision_model(pixel_values=pixels).pooler_output
    )
    theirs_text = clip_hf.text_projection(clip_hf.text_model(**batch).pooler_output)
print(
    "image vectors, largest difference:", (ours_image - theirs_image).abs().max().item()
)
print("text vectors, largest difference:", (ours_text - theirs_text).abs().max().item())
print("logit scale:", clip.logit_scale.item(), "->", clip.logit_scale.exp().item())

batch_ids = BATCH
batch_photos = [load(coco_thumbs[i]) for i in batch_ids]
batch_texts = [coco_captions[i][0] for i in batch_ids]
with torch.no_grad():
    pixels = torch.stack([clip_pixels(photo) for photo in batch_photos]).to(device)
    tokens = tokenizer(
        batch_texts, padding=True, truncation=True, return_tensors="pt"
    ).to(device)
    u = F.normalize(clip.image(pixels), dim=-1).cpu()
    v = F.normalize(
        clip.text(tokens["input_ids"], tokens["attention_mask"]), dim=-1
    ).cpu()
cosine = u @ v.T
scale = clip.logit_scale.exp().item()
logits = scale * cosine
labels = torch.arange(len(batch_ids))
loss_photos = F.cross_entropy(logits, labels)
loss_captions = F.cross_entropy(logits.T, labels)
print("1 / temperature:", round(scale, 1))
print("photos find their captions:", round(loss_photos.item(), 3))
print("captions find their photos:", round(loss_captions.item(), 3))
print("CLIP's loss:", round((loss_photos.item() + loss_captions.item()) / 2, 3))

fig, ax = plt.subplots(figsize=(9.5, 8.4))
ax.imshow(cosine, cmap=heat, vmin=cosine.min(), vmax=cosine.max())
for i in range(len(batch_ids)):
    for j in range(len(batch_ids)):
        strong = cosine[i, j] > cosine.min() + 0.7 * (cosine.max() - cosine.min())
        ink = "white" if strong else COLORS["ink"]
        ax.text(
            j, i, f"{cosine[i, j]:.2f}", ha="center", va="center", fontsize=9, color=ink
        )
short = [textwrap.shorten(t, 28, placeholder="…") for t in batch_texts]
ax.set_xticks(range(len(short)), short, rotation=40, ha="right", fontsize=8)
ax.set_yticks([])
ax.grid(False)
for i, row_photo in enumerate(batch_photos):  # each photo as the label of its row
    small = row_photo.resize((round(80 * row_photo.width / row_photo.height), 80))
    picture = OffsetImage(np.asarray(small), zoom=0.55)
    label = AnnotationBbox(
        picture, (-0.6, i), box_alignment=(1, 0.5), frameon=False, annotation_clip=False
    )
    ax.add_artist(label)
ax.set_title(
    "cosine of every photo with every caption; the diagonal holds the true pairs",
    loc="left",
)
plt.show()

row = 1
fig, axes = plt.subplots(1, 3, figsize=(12, 3.2), sharey=True)
for ax, tau in zip(axes, [1.0, 0.07, 1 / scale]):
    p = (cosine[row] / tau).softmax(-1)
    colors = [COLORS["dec"] if j == row else COLORS["line"] for j in range(len(p))]
    ax.bar(range(len(p)), p, color=colors)
    ax.set_title(f"τ = {tau:.3g}", fontsize=11)
    ax.set_xticks(range(len(p)), [f"t{j + 1}" for j in range(len(p))], fontsize=8)
axes[0].set_ylabel("probability of each caption")
plt.show()

@torch.no_grad()
def embed_photos(thumbs, pixels_of, image_tower, batch_size=256):
    vectors = []
    with ThreadPoolExecutor(8) as pool:  # decode and resize in parallel
        for start in range(0, len(thumbs), batch_size):
            chunk = thumbs[start : start + batch_size]
            pixels = torch.stack(list(pool.map(lambda t: pixels_of(load(t)), chunk)))
            with torch.autocast(device, dtype=torch.float16, enabled=device == "cuda"):
                out = image_tower(pixels.to(device))
            vectors.append(F.normalize(out.float(), dim=-1).cpu())
    return torch.cat(vectors)


@torch.no_grad()
def clip_texts(texts, batch_size=512):
    vectors = []
    for start in range(0, len(texts), batch_size):
        chunk = list(texts[start : start + batch_size])
        tokens = tokenizer(chunk, padding=True, truncation=True, return_tensors="pt")
        tokens = tokens.to(device)
        with torch.autocast(device, dtype=torch.float16, enabled=device == "cuda"):
            out = clip.text(tokens["input_ids"], tokens["attention_mask"])
        vectors.append(F.normalize(out.float(), dim=-1).cpu())
    return torch.cat(vectors)


start = time.time()
clip_geo_photos = embed_photos(geo_thumbs, clip_pixels, clip.image)
clip_coco_photos = embed_photos(coco_thumbs, clip_pixels, clip.image)
print(
    "Country211:", tuple(clip_geo_photos.shape), "COCO:", tuple(clip_coco_photos.shape)
)
print(f"in {time.time() - start:.0f} s")

def country_classifier(templates, encode=clip_texts):
    rows = []
    for country in COUNTRIES:
        sentences = [template.format(c=country) for template in templates]
        mean = encode(sentences).mean(dim=0)  # several sentences per class: average
        rows.append(F.normalize(mean, dim=-1))
    return torch.stack(rows)


def accuracy(scores, truth):
    top5 = scores.topk(5, dim=-1).indices
    top1 = (top5[:, 0] == truth).float().mean().item()
    within5 = (top5 == truth[:, None]).any(dim=-1).float().mean().item()
    return top1, within5

TEMPLATES = {
    "Italy": ["{c}"],
    "a photo of Italy.": ["a photo of {c}."],
    "a photo i took in Italy.": ["a photo i took in {c}."],
    "the paper's five sentences, averaged": PAPER_TEMPLATES,
}
rows = []
for name, templates in TEMPLATES.items():
    scores = clip_geo_photos @ country_classifier(templates).T
    top1, within5 = accuracy(scores, geo_truth)
    rows.append({"classes written as": name, "top-1": top1, "in the top 5": within5})
print("chance:", round(1 / len(COUNTRIES), 4))
table(pd.DataFrame(rows), digits=3)

clip_geo = country_classifier(PAPER_TEMPLATES)
predicted = (clip_geo_photos @ clip_geo.T).argmax(dim=-1)
present = sorted(set(geo_truth.tolist()))
per_country = {
    k: (predicted[geo_truth == k] == k).float().mean().item() for k in present
}
top = sorted(present, key=lambda k: per_country[k])[-20:]
fig, ax = plt.subplots(figsize=(8, 6.4))
ax.barh([COUNTRIES[k] for k in top], [per_country[k] for k in top], color=COLORS["enc"])
ax.axvline(1 / len(COUNTRIES), color=COLORS["muted"], linestyle="--", linewidth=1)
ax.set_xlabel("share of the country's photos placed correctly")
ax.set_title("the twenty countries CLIP ViT-B/32 places best")
ax.grid(axis="x")
ax.grid(axis="y", visible=False)
plt.show()
never = sum(per_country[k] == 0 for k in present)
print(f"countries where not one photo is placed correctly: {never} of {len(present)}")

mistakes = Counter()
for truth, guess in zip(geo_truth.tolist(), predicted.tolist()):
    if truth != guess:
        mistakes[(COUNTRIES[truth], COUNTRIES[guess])] += 1
rows = []
for (truth, guess), k in mistakes.most_common(10):
    rows.append({"taken in": truth, "placed in": guess, "photos": k})
display(table(pd.DataFrame(rows), digits=0))
named = Counter(COUNTRIES[guess] for guess in predicted.tolist())
rows = [{"country": c, "photos placed there": k} for c, k in named.most_common(8)]
table(pd.DataFrame(rows), digits=0)

QUERIES = [
    "a cat sleeping on a laptop",
    "a giraffe eating leaves from a tree",
    "an elephant drinking water",
    "three zebras and one giraffe",
]
query_vectors = clip_texts(QUERIES)
best = (query_vectors @ clip_coco_photos.T).topk(5, dim=-1)
fig, axes = plt.subplots(len(QUERIES), 5, figsize=(12, 2.6 * len(QUERIES)))
for row_axes, query, scores, found in zip(axes, QUERIES, best.values, best.indices):
    for ax, score, i in zip(row_axes, scores.tolist(), found.tolist()):
        ax.imshow(load(coco_thumbs[i]))
        ax.set_title(f"{score:.3f}", fontsize=8)
        ax.axis("off")
    row_axes[0].text(
        -0.08,
        0.5,
        query,
        transform=row_axes[0].transAxes,
        rotation=90,
        ha="right",
        va="center",
        fontsize=9,
    )
plt.tight_layout()
plt.show()

def summary(ranks):
    result = {f"recall@{k}": (ranks < k).float().mean().item() for k in (1, 5, 10)}
    result["median rank"] = ranks.float().median().item() + 1
    return result


def text_to_image(caption_vectors, photo_vectors, owner, chunk=4096):
    ranks = []
    for start in range(0, len(caption_vectors), chunk):
        scores = caption_vectors[start : start + chunk] @ photo_vectors.T
        right = scores[torch.arange(len(scores)), owner[start : start + chunk]]
        ranks.append(
            (scores > right[:, None]).sum(dim=-1)
        )  # photos above the right one
    return torch.cat(ranks)


def image_to_text(photo_vectors, caption_vectors, owner, chunk=1024):
    ranks = []
    for start in range(0, len(photo_vectors), chunk):
        scores = photo_vectors[start : start + chunk] @ caption_vectors.T
        mine = owner[None, :] == torch.arange(start, start + len(scores))[:, None]
        best_own = scores.masked_fill(~mine, float("-inf")).max(dim=-1).values
        ranks.append((scores > best_own[:, None]).sum(dim=-1))
    return torch.cat(ranks)


clip_captions = clip_texts(captions)
clip_t2i = summary(text_to_image(clip_captions, clip_coco_photos, owner))
clip_i2t = summary(image_to_text(clip_coco_photos, clip_captions, owner))
search_rows = [
    {"model": "CLIP ViT-B/32", "search": "caption → photo", **clip_t2i},
    {"model": "CLIP ViT-B/32", "search": "photo → caption", **clip_i2t},
]
print(len(captions), "captions,", len(coco_thumbs), "photos")
table(pd.DataFrame(search_rows), digits=3)

pair_photos = clip_coco_photos
pair_texts = clip_captions[first]
mixed = torch.randperm(len(pair_texts), generator=torch.Generator().manual_seed(SEED))
rows = [
    {
        "pair": "a photo and its caption",
        "mean cosine": (pair_photos * pair_texts).sum(-1).mean().item(),
    },
    {
        "pair": "a photo and a random caption",
        "mean cosine": (pair_photos * pair_texts[mixed]).sum(-1).mean().item(),
    },
    {
        "pair": "two random photos",
        "mean cosine": (pair_photos @ pair_photos.T).mean().item(),
    },
    {
        "pair": "two random captions",
        "mean cosine": (pair_texts @ pair_texts.T).mean().item(),
    },
]
table(pd.DataFrame(rows), digits=3)

both = torch.cat([pair_photos[:1000], pair_texts[:1000]])
centered = both - both.mean(dim=0)
_, _, directions = torch.linalg.svd(centered, full_matrices=False)
flat = centered @ directions[:2].T
k = len(flat) // 2
fig, ax = plt.subplots(figsize=(6.4, 5))
ax.scatter(
    flat[:k, 0], flat[:k, 1], s=6, color=COLORS["enc"], alpha=0.6, label="photos"
)
ax.scatter(
    flat[k:, 0], flat[k:, 1], s=6, color=COLORS["data"], alpha=0.6, label="captions"
)
for i in range(0, k, 100):
    ax.plot(
        [flat[i, 0], flat[k + i, 0]],
        [flat[i, 1], flat[k + i, 1]],
        color=COLORS["line"],
        linewidth=0.6,
    )
ax.set_xlabel("first principal direction")
ax.set_ylabel("second")
ax.set_title("photos and their captions in CLIP's space, projected to 2D")
ax.legend(frameon=False)
ax.grid(False)
plt.show()

offset = pair_photos.mean(dim=0) - pair_texts.mean(dim=0)
print("distance between the two centres:", round(offset.norm().item(), 3))
moved = F.normalize(clip_captions + offset, dim=-1)
search_rows.append(
    {
        "model": "CLIP, captions moved across the gap",
        "search": "caption → photo",
        **summary(text_to_image(moved, clip_coco_photos, owner)),
    }
)
table(pd.DataFrame(search_rows), digits=3)

FONT = font_manager.findfont(
    font_manager.FontProperties(family="DejaVu Sans", weight="bold")
)


def write_on(photo, word, height):
    photo = photo.copy()
    draw = ImageDraw.Draw(photo)
    font = ImageFont.truetype(FONT, max(8, int(photo.height * height)))
    left, top, right, bottom = draw.textbbox((0, 0), word, font=font)
    x = (photo.width - (right - left)) // 2
    y = photo.height - (bottom - top) - int(photo.height * 0.08)
    draw.rectangle(
        (x - 6, y + top - 6, x + right - left + 6, y + bottom + 6), fill="white"
    )
    draw.text((x, y), word, fill="black", font=font)
    return photo


japan = COUNTRIES.index("Japan")


def placed_right(guesses):
    # the photos are sorted by country: shuffle, so that the 200 come from many countries
    order = torch.randperm(
        len(geo_thumbs), generator=torch.Generator().manual_seed(SEED)
    )
    right = []
    for position in order.tolist():
        if guesses[position] == geo_truth[position] != japan:
            right.append(position)
    return right[:200]


right_ones = placed_right(predicted)
print(len(right_ones), "photos that CLIP placed correctly, none of them taken in Japan")

SIZES = [0.04, 0.07, 0.10, 0.15]


def after_writing(word, height, pixels_of, image_tower, classifier, chosen):
    stamped = [write_on(load(geo_thumbs[p]), word, height) for p in chosen]
    with torch.no_grad():
        pixels = torch.stack([pixels_of(photo) for photo in stamped]).to(device)
        vectors = F.normalize(image_tower(pixels).float(), dim=-1).cpu()
    guesses = (vectors @ classifier.T).argmax(dim=-1)
    return {
        "now placed in Japan": (guesses == japan).float().mean().item(),
        "still placed correctly": (guesses == geo_truth[chosen]).float().mean().item(),
    }


def attack(pixels_of, image_tower, classifier, chosen):
    rows = []
    for word, heights in [("JAPAN", SIZES), ("PHOTO", [0.10])]:
        for height in heights:
            moved = after_writing(
                word, height, pixels_of, image_tower, classifier, chosen
            )
            rows.append({"word": word, "text height": height, **moved})
    return pd.DataFrame(rows)


table(attack(clip_pixels, clip.image, clip_geo, right_ones), digits=3)

fig, axes = plt.subplots(2, 4, figsize=(12, 6.2))
for column, position in enumerate(right_ones[:4]):
    original = load(geo_thumbs[position])
    for row, image in enumerate([original, write_on(original, "JAPAN", 0.10)]):
        with torch.no_grad():
            vector = F.normalize(
                clip.image(clip_pixels(image)[None].to(device)).float(), dim=-1
            ).cpu()
        probabilities = (scale * vector @ clip_geo.T).softmax(dim=-1)[0]
        guess = probabilities.argmax().item()
        axes[row, column].imshow(image)
        axes[row, column].set_title(
            f"{COUNTRIES[guess]}, {probabilities[guess]:.0%}", fontsize=9
        )
        axes[row, column].axis("off")
plt.tight_layout()
plt.show()