create_reference_notebook.py

"""
Create clean reference solution notebook for Visual Genomes.
This is NOT given to contestants. It's for organizers only.
"""
import nbformat

nb = nbformat.v4.new_notebook()

nb.cells.append(nbformat.v4.new_markdown_cell("""# Visual Genomes - Reference Solutions (Organizer Only)

This notebook demonstrates 4 progressively sophisticated approaches.

| Approach | Combined Score | Key Technique |
|----------|:--:|-------------|
| Baseline | 0.123 | Most common predicate for all pairs |
| V1 | 0.167 | Self-attention weight ranking |
| V2 | 0.301 | Hidden state concatenation + LogReg |
| V3 | 0.433 | Multi-layer features + MLP |
| V4 | 0.439 | Label smoothing + connectivity |
"""))

nb.cells.append(nbformat.v4.new_code_cell("""import os, json, random, numpy as np, torch, torch.nn as nn, pandas as pd
from collections import Counter, defaultdict
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler
from torch.utils.data import Dataset, DataLoader

random.seed(2026); np.random.seed(2026); torch.manual_seed(2026)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

BASE = '/kaggle/input/visual-genomes-aicc'
train_data = torch.load(f'{BASE}/train.pt', weights_only=False)
val_data = torch.load(f'{BASE}/val.pt', weights_only=False)
test_data = torch.load(f'{BASE}/test.pt', weights_only=False)

with open(f'{BASE}/predicates.json', 'r') as f:
    pred_info = json.load(f)
predicates = pred_info['predicates']
n_pred = len(predicates)
print(f"Train: {len(train_data)}, Val: {len(val_data)}, Test: {len(test_data)}, Predicates: {n_pred}")
"""))

nb.cells.append(nbformat.v4.new_code_cell("""# Evaluation helpers
def compute_recall_at_k(predictions, ground_truth, k):
    recalls = []
    for img_id, gt_triplets in ground_truth.items():
        if not gt_triplets: continue
        preds = predictions.get(img_id, [])[:k]
        pred_set = set((p['subject_idx'], p['object_idx'], p['predicate']) for p in preds)
        recalls.append(len(gt_triplets & pred_set) / len(gt_triplets))
    return np.mean(recalls) if recalls else 0.0

def compute_mean_recall_at_k(predictions, ground_truth, k):
    per_class_hits, per_class_total = defaultdict(int), defaultdict(int)
    for img_id, gt_triplets in ground_truth.items():
        preds = predictions.get(img_id, [])[:k]
        pred_set = set((p['subject_idx'], p['object_idx'], p['predicate']) for p in preds)
        for s, o, p in gt_triplets:
            per_class_total[p] += 1
            if (s, o, p) in pred_set: per_class_hits[p] += 1
    return np.mean([per_class_hits[p]/per_class_total[p] for p in per_class_total]) if per_class_total else 0.0

def evaluate(rows, data):
    gt = {e['image_id']: set(tuple(t) for t in e['triplets']) for e in data}
    preds = defaultdict(list)
    for r in rows: preds[r['image_id']].append(r)
    for k in preds: preds[k].sort(key=lambda x: x['confidence'], reverse=True)
    r20 = compute_recall_at_k(preds, gt, 20)
    r50 = compute_recall_at_k(preds, gt, 50)
    mr20 = compute_mean_recall_at_k(preds, gt, 20)
    combined = 0.4*r20 + 0.4*r50 + 0.2*mr20
    print(f"R@20: {r20:.4f}  R@50: {r50:.4f}  mR@20: {mr20:.4f}  Combined: {combined:.4f}")
    return combined

val_gt = val_data  # val has triplets
"""))

# V1: Attention thresholding
nb.cells.append(nbformat.v4.new_markdown_cell("## V1: Attention Weight Ranking"))
nb.cells.append(nbformat.v4.new_code_cell("""pred_freq = Counter()
for e in train_data:
    for _,_,p in e['triplets']: pred_freq[p] += 1
top5 = [p for p,_ in pred_freq.most_common(5)]

rows_v1 = []
for entry in val_data:
    attn = entry['self_attn_weights'][-1].float().mean(0)
    pairs = [(i,j,attn[i,j].item()) for i in range(entry['n_det']) for j in range(entry['n_det']) if i!=j]
    pairs.sort(key=lambda x: x[2], reverse=True)
    for si,oi,a in pairs:
        for r,p in enumerate(top5):
            rows_v1.append({'image_id': entry['image_id'], 'subject_idx': si, 'object_idx': oi,
                           'predicate': int(p), 'confidence': a*(1-0.01*r)})
print("V1 Val:"); evaluate(rows_v1, val_gt)
"""))

# V2: Hidden concat + LogReg
nb.cells.append(nbformat.v4.new_markdown_cell("## V2: Hidden State Concatenation + LogReg"))
nb.cells.append(nbformat.v4.new_code_cell("""X_train, y_train = [], []
for entry in train_data:
    h = entry['decoder_hidden'][-1].float()
    gt_pairs = set()
    for s,o,p in entry['triplets']:
        if s < entry['n_det'] and o < entry['n_det']:
            X_train.append(torch.cat([h[s], h[o]]).numpy()); y_train.append(p); gt_pairs.add((s,o))
    neg = [(i,j) for i in range(entry['n_det']) for j in range(entry['n_det']) if i!=j and (i,j) not in gt_pairs]
    for s,o in random.sample(neg, min(len(neg), len(gt_pairs)*3)):
        X_train.append(torch.cat([h[s], h[o]]).numpy()); y_train.append(n_pred)

X_train, y_train = np.array(X_train), np.array(y_train)
scaler = StandardScaler(); X_s = scaler.fit_transform(X_train)
clf = LogisticRegression(max_iter=1000, C=1.0, solver='lbfgs', random_state=2026)
clf.fit(X_s, y_train); print(f"Train acc: {clf.score(X_s, y_train):.4f}")

rows_v2 = []
for entry in val_data:
    h = entry['decoder_hidden'][-1].float()
    feats, pairs = [], []
    for i in range(entry['n_det']):
        for j in range(entry['n_det']):
            if i!=j: feats.append(torch.cat([h[i],h[j]]).numpy()); pairs.append((i,j))
    if not feats: continue
    probs = clf.predict_proba(scaler.transform(np.array(feats)))
    for idx,(si,oi) in enumerate(pairs):
        for ci,cls in enumerate(clf.classes_):
            if cls != n_pred and probs[idx,ci] > 0.01:
                rows_v2.append({'image_id': entry['image_id'], 'subject_idx': si, 'object_idx': oi,
                               'predicate': int(cls), 'confidence': float(probs[idx,ci])})
print("V2 Val:"); evaluate(rows_v2, val_gt)
"""))

# V3: Multi-layer MLP
nb.cells.append(nbformat.v4.new_markdown_cell("## V3: Multi-Layer Features + MLP"))
nb.cells.append(nbformat.v4.new_code_cell("""class RelDS(Dataset):
    def __init__(self, entries, n_pred, train=True, neg_r=3):
        self.X, self.Y = [], []
        for e in entries:
            nd = e['n_det']; lab = torch.zeros(nd,nd,n_pred); gt = set()
            for s,o,p in e['triplets']:
                if s<nd and o<nd: lab[s,o,p]=1.0; gt.add((s,o))
            pf = {}
            for i in range(nd):
                for j in range(nd):
                    if i==j: continue
                    parts = []
                    for h in e['decoder_hidden']: hf=h.float(); parts.extend([hf[i],hf[j]])
                    for a in e['self_attn_weights']: af=a.float(); parts.extend([af[:,i,j],af[:,j,i]])
                    pf[(i,j)] = torch.cat(parts)
            for s,o in gt:
                if (s,o) in pf: self.X.append(pf[(s,o)]); self.Y.append(lab[s,o])
            if train:
                neg = [(i,j) for i in range(nd) for j in range(nd) if i!=j and (i,j) not in gt and (i,j) in pf]
                for s,o in random.sample(neg, min(len(neg),len(gt)*neg_r)):
                    self.X.append(pf[(s,o)]); self.Y.append(lab[s,o])
    def __len__(self): return len(self.X)
    def __getitem__(self,i): return self.X[i], self.Y[i]

class MLP(nn.Module):
    def __init__(self, d_in, n_p, hid=256):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(d_in,hid),nn.ReLU(),nn.Dropout(0.3),
                                 nn.Linear(hid,hid//2),nn.ReLU(),nn.Dropout(0.2),
                                 nn.Linear(hid//2,n_p))
    def forward(self,x): return self.net(x)

tds = RelDS(train_data, n_pred, True, 3); vds = RelDS(val_data, n_pred, False, 5)
d_in = tds[0][0].shape[0]; print(f"Train: {len(tds)}, Val: {len(vds)}, Dim: {d_in}")
model = MLP(d_in, n_pred).to(device)

pc = torch.zeros(n_pred)
for _,y in tds: pc += y
pw = ((len(tds)-pc)/(pc+1)).clamp(max=50).to(device)
crit = nn.BCEWithLogitsLoss(pos_weight=pw)
opt = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-4)
sch = torch.optim.lr_scheduler.CosineAnnealingLR(opt, 30)
loader = DataLoader(tds, 256, shuffle=True)
best_vl, best_st = float('inf'), None

for ep in range(30):
    model.train()
    for f,y in loader:
        f,y = f.to(device),y.to(device); opt.zero_grad()
        crit(model(f),y).backward(); opt.step()
    sch.step()
    model.eval(); vl = []
    with torch.no_grad():
        for f,y in DataLoader(vds,512): f,y=f.to(device),y.to(device); vl.append(crit(model(f),y).item())
    avg = np.mean(vl)
    if avg < best_vl: best_vl=avg; best_st={k:v.cpu().clone() for k,v in model.state_dict().items()}
    if (ep+1)%10==0: print(f"Epoch {ep+1}: val_loss={avg:.4f}")

model.load_state_dict(best_st); model.to(device).eval()

rows_v3 = []
for entry in val_data:
    for i in range(entry['n_det']):
        for j in range(entry['n_det']):
            if i==j: continue
            parts = []
            for h in entry['decoder_hidden']: hf=h.float(); parts.extend([hf[i],hf[j]])
            for a in entry['self_attn_weights']: af=a.float(); parts.extend([af[:,i,j],af[:,j,i]])
            feat = torch.cat(parts).unsqueeze(0).to(device)
            with torch.no_grad(): probs = torch.sigmoid(model(feat)[0]).cpu()
            for p in range(n_pred):
                if probs[p].item() > 0.05:
                    rows_v3.append({'image_id': entry['image_id'], 'subject_idx': i,
                                   'object_idx': j, 'predicate': p, 'confidence': probs[p].item()})
print("V3 Val:"); evaluate(rows_v3, val_gt)
"""))

# V4: Smoothing + Connectivity
nb.cells.append(nbformat.v4.new_markdown_cell("## V4: Label Smoothing + Connectivity"))
nb.cells.append(nbformat.v4.new_code_cell("""def build_feat(e, i, j):
    parts = []
    for h in e['decoder_hidden']: hf=h.float(); parts.extend([hf[i],hf[j]])
    for a in e['self_attn_weights']: af=a.float(); parts.extend([af[:,i,j],af[:,j,i]])
    parts.append(e['det_scores'].float()[[i,j]])
    return torch.cat(parts)

class SmDS(Dataset):
    def __init__(self, entries, n_pred, train=True, neg_r=3, alpha=0.1):
        self.X, self.RL, self.CL = [], [], []
        for e in entries:
            nd=e['n_det']; sc=e['det_scores'].float()
            rl=torch.zeros(nd,nd,n_pred); cl=torch.zeros(nd,nd); gt=set()
            for s,o,p in e['triplets']:
                if s<nd and o<nd:
                    us=1-sc[s].item(); uo=1-sc[o].item()
                    rl[s,o,p]=(1-max(us,alpha))*(1-max(uo,alpha)); cl[s,o]=1.0; gt.add((s,o))
            for s,o in gt:
                self.X.append(build_feat(e,s,o)); self.RL.append(rl[s,o]); self.CL.append(torch.tensor(1.0))
            if train:
                neg=[(i,j) for i in range(nd) for j in range(nd) if i!=j and (i,j) not in gt]
                la=e['self_attn_weights'][-1].float().mean(0)
                ns=sorted([(i,j,la[i,j].item()) for i,j in neg], key=lambda x:x[2], reverse=True)
                nn_=min(len(neg),len(gt)*neg_r); nh=nn_//2
                for i,j,_ in ns[:nh]:
                    self.X.append(build_feat(e,i,j)); self.RL.append(torch.zeros(n_pred)); self.CL.append(torch.tensor(0.0))
                for i,j in random.sample(neg, min(nn_-nh,len(neg))):
                    self.X.append(build_feat(e,i,j)); self.RL.append(torch.zeros(n_pred)); self.CL.append(torch.tensor(0.0))
    def __len__(self): return len(self.X)
    def __getitem__(self,i): return self.X[i], self.RL[i], self.CL[i]

class DH(nn.Module):
    def __init__(self,d_in,n_p,hid=256):
        super().__init__()
        self.bb=nn.Sequential(nn.Linear(d_in,hid),nn.ReLU(),nn.Dropout(0.3),nn.Linear(hid,hid//2),nn.ReLU(),nn.Dropout(0.2))
        self.rh=nn.Linear(hid//2,n_p); self.ch=nn.Linear(hid//2,1)
    def forward(self,x): h=self.bb(x); return self.rh(h), self.ch(h).squeeze(-1)

tds4=SmDS(train_data,n_pred,True); vds4=SmDS(val_data,n_pred,False)
d4=tds4[0][0].shape[0]; print(f"Train: {len(tds4)}, Val: {len(vds4)}, Dim: {d4}")
m4=DH(d4,n_pred).to(device)

pc4=torch.zeros(n_pred)
for _,r,_ in tds4: pc4+=(r>0).float()
pw4=((len(tds4)-pc4)/(pc4+1)).clamp(max=50).to(device)
rc=nn.BCEWithLogitsLoss(pos_weight=pw4); cc=nn.BCEWithLogitsLoss()
o4=torch.optim.Adam(m4.parameters(),lr=1e-3,weight_decay=1e-4)
s4=torch.optim.lr_scheduler.CosineAnnealingLR(o4,40)
ld4=DataLoader(tds4,256,shuffle=True); bv4,bs4=float('inf'),None

for ep in range(40):
    m4.train()
    for f,r,c in ld4:
        f,r,c=f.to(device),r.to(device),c.to(device); o4.zero_grad()
        rl,cl=m4(f); (rc(rl,r)+0.5*cc(cl,c)).backward(); o4.step()
    s4.step()
    m4.eval(); vl=[]
    with torch.no_grad():
        for f,r,c in DataLoader(vds4,512):
            f,r,c=f.to(device),r.to(device),c.to(device)
            rl,cl=m4(f); vl.append((rc(rl,r)+0.5*cc(cl,c)).item())
    av=np.mean(vl)
    if av<bv4: bv4=av; bs4={k:v.cpu().clone() for k,v in m4.state_dict().items()}
    if (ep+1)%10==0: print(f"Epoch {ep+1}: val={av:.4f}")

m4.load_state_dict(bs4); m4.to(device).eval()

# Generate test submission with V4
rows_final = []
for entry in test_data:
    for i in range(entry['n_det']):
        for j in range(entry['n_det']):
            if i==j: continue
            feat = build_feat(entry,i,j).unsqueeze(0).to(device)
            with torch.no_grad():
                rl,cl = m4(feat)
                rp = torch.sigmoid(rl[0]).cpu()
                cp = torch.sigmoid(cl[0]).cpu().item()
            for p in range(n_pred):
                conf = rp[p].item() * cp
                if conf > 0.02:
                    rows_final.append({'image_id': entry['image_id'], 'subject_idx': i,
                                      'object_idx': j, 'predicate': p, 'confidence': conf})

submission = pd.DataFrame(rows_final)
submission.to_csv('submission.csv', index=False)
print(f"Final submission: {len(submission)} predictions")
"""))

nbformat.write(nb, 'reference_solution.ipynb')
print("Created reference_solution.ipynb")