File size: 2,404 Bytes
723126c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 | import json
from pathlib import Path
import numpy as np
import torch
import torch.nn.functional as F
from transformers import CLIPTextModelWithProjection, CLIPTokenizer
EVALS = Path("/home/s56fasla/MOTIONX/evals")
CLIP_NAME = "openai/clip-vit-base-patch32"
ARMS = ["base", "baseline", "ski"]
def lcs(a, b):
m = [[0] * (len(b) + 1) for _ in range(len(a) + 1)]
for i in range(len(a)):
for j in range(len(b)):
m[i + 1][j + 1] = m[i][j] + 1 if a[i] == b[j] else max(m[i][j + 1], m[i + 1][j])
return m[-1][-1]
def rouge_l(pred, ref):
p, r = pred.lower().split(), ref.lower().split()
if not p or not r:
return 0.0
l = lcs(p, r)
prec, rec = l / len(p), l / len(r)
return 0.0 if prec + rec == 0 else 2 * prec * rec / (prec + rec)
dev = "cuda" if torch.cuda.is_available() else "cpu"
tok = CLIPTokenizer.from_pretrained(CLIP_NAME)
enc = CLIPTextModelWithProjection.from_pretrained(CLIP_NAME).to(dev).eval()
def embed(texts):
out = []
with torch.no_grad():
for i in range(0, len(texts), 256):
b = tok(texts[i:i + 256], padding=True, truncation=True,
max_length=77, return_tensors="pt").to(dev)
out.append(F.normalize(enc(**b).text_embeds, dim=-1).cpu())
return torch.cat(out).float()
results = {}
for arm in ARMS:
f = EVALS / f"preds_{arm}.json"
if not f.exists():
print("missing", f, flush=True)
continue
preds = json.load(open(f))
P = embed([p["pred"] for p in preds])
R = embed([p["ref"] for p in preds])
cos = float((P * R).sum(-1).mean())
sim = P @ R.T
tgt = torch.arange(len(preds))
rank = (sim > sim[tgt, tgt].unsqueeze(1)).sum(1)
rl = float(np.mean([rouge_l(p["pred"], p["ref"]) for p in preds]))
results[arm] = {
"n": len(preds),
"clip_cos": round(cos, 4),
"rouge_l": round(rl, 4),
"retr_r1": round(float((rank < 1).float().mean()), 4),
"retr_r5": round(float((rank < 5).float().mean()), 4),
"medr": float(rank.float().median()) + 1,
"mean_words": round(float(np.mean([len(p["pred"].split()) for p in preds])), 1),
"chance_r1": round(1.0 / len(preds), 5),
}
print(arm, json.dumps(results[arm]), flush=True)
json.dump(results, open(EVALS / "caption_scores.json", "w"), indent=1)
print("saved", EVALS / "caption_scores.json", flush=True)
|