1import sys, json, importlib.util
 2from pathlib import Path
 3import numpy as np
 4import torch
 5import torch.nn as nn
 6import torch.nn.functional as F
 7sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
 8from bench import train_model, evaluate, sweep_baseline, make_report, count_params
 9TRACK='multitoken_diffusion'; DIM=8
10_spec=importlib.util.spec_from_file_location('registered_multitoken_diffusion','/home/maxwelhelp/all/math2nn/bench/custom_tracks/multitoken_diffusion.py')
11_mod=importlib.util.module_from_spec(_spec); _spec.loader.exec_module(_mod)
12class LocalBaseline(nn.Module):
13    def __init__(self, hidden=32, radius=1):
14        super().__init__(); k=2*radius+1
15        self.body=nn.Sequential(nn.Conv1d(2,hidden,k,padding=radius),nn.ReLU(),nn.Conv1d(hidden,hidden,k,padding=radius),nn.ReLU(),nn.Conv1d(hidden,1,1))
16    def forward(self,x):
17        z=x[:,:DIM].unsqueeze(1); t=x[:,DIM:].unsqueeze(1).expand(-1,1,DIM)
18        return self.body(torch.cat([z,t],1)).squeeze(1)
19class RGPyramid(nn.Module):
20    def __init__(self, hidden=32, radius=1):
21        super().__init__(); k=2*radius+1
22        self.fine=nn.Sequential(nn.Conv1d(2,hidden,k,padding=radius),nn.ReLU(),nn.Conv1d(hidden,hidden,k,padding=radius),nn.ReLU())
23        self.coarse=nn.Sequential(nn.Conv1d(2,hidden,k,padding=radius),nn.ReLU(),nn.Conv1d(hidden,hidden,1),nn.ReLU())
24        self.out=nn.Conv1d(2*hidden,1,1)
25    def forward(self,x):
26        z=x[:,:DIM].unsqueeze(1); t=x[:,DIM:].unsqueeze(1).expand(-1,1,DIM)
27        coarse=F.avg_pool1d(torch.cat([z,t],1),2,2); c=F.interpolate(self.coarse(coarse),size=DIM,mode='linear',align_corners=False)
28        f=self.fine(torch.cat([z,t],1)); return self.out(torch.cat([f,c],1)).squeeze(1)
29def dataset(seed): return _mod.get_dataset(seed,n_train=400,n_test=200)
30def run(kind,seed,cfg,detail=False):
31    torch.manual_seed(seed); np.random.seed(seed); ds=dataset(seed)
32    model=LocalBaseline(cfg['hidden'],1) if kind=='baseline' else RGPyramid(cfg['hidden'],1)
33    model,train_metric,_=train_model(model,ds,epochs=cfg['epochs'],lr=cfg['lr'],batch=128,log=lambda *a,**k:None)
34    if model is None: raise RuntimeError('bench training returned None')
35    # Evaluation on CPU avoids a shared-GPU Conv1d engine failure; both systems use this identically.
36    model=model.cpu().eval()
37    with torch.no_grad(): pred=model(ds['xte']).numpy()
38    metric=float(np.mean((pred-ds['yte'].numpy())**2))
39    if not detail:return metric
40    x=ds['xte'][:128].clone(); x2=x.clone(); x2[:,0]+=0.5
41    with torch.no_grad(): delta=np.abs(model(x2).numpy()-model(x).numpy())
42    return metric,{'params':count_params(model),'mean_output_sensitivity':float(delta.mean()),'distant_sensitivity_0_to_7':float(delta[:,7].mean()),'near_sensitivity_0_to_1':float(delta[:,1].mean()),'train_metric':float(train_metric)}
43def factory(kind,cfg): return lambda seed: run(kind,seed,cfg)
44def main():
45    grid=[{'lr':1e-3,'epochs':12,'hidden':32},{'lr':3e-3,'epochs':12,'hidden':32},{'lr':6e-3,'epochs':12,'hidden':32}]
46    base=sweep_baseline(lambda c:factory('baseline',c),grid,seeds=(0,1,2,3)); bcfg=base['best_cfg']
47    base_full=evaluate(factory('baseline',bcfg),seeds=tuple(range(8)))
48    idea_runs=[(c,evaluate(factory('idea',c),seeds=tuple(range(8)))) for c in grid]
49    best_cfg,best=min(idea_runs,key=lambda q:q[1]['mean'])
50    isig=[run('idea',s,best_cfg,True)[1] for s in range(8)]; bsig=[run('baseline',s,bcfg,True)[1] for s in range(8)]
51    im=float(np.mean([q['distant_sensitivity_0_to_7'] for q in isig])); bm=float(np.mean([q['distant_sensitivity_0_to_7'] for q in bsig]))
52    signature={'prediction':'coarse pooled conditioning transmits distant-token influence beyond a radius-1 local field','predicted_vs_observed':{'predicted':'idea distant sensitivity > baseline','observed_idea_mean_distant_sensitivity':im,'observed_baseline_mean_distant_sensitivity':bm},'trained_model_observations':isig,'confirmed':bool(im>bm)}
53    report=make_report(TRACK,'custom_local_vs_rg_pyramid',{'best_cfg':bcfg,'sweep':base['sweep'],'full':base_full},best,signature)
54    report['idea_sweep']=[{'cfg':c,'result':r} for c,r in idea_runs]
55    report['custom_track']={'name':TRACK,'file':'/home/maxwelhelp/all/math2nn/bench/custom_tracks/multitoken_diffusion.py','domain':'diffusion-sampling'}
56    report['architecture']={'baseline':'local radius-1 convolutional velocity field','idea':'dyadic pooled coarse field plus local fine residual conditioning','baseline_params':count_params(LocalBaseline()),'idea_params':count_params(RGPyramid())}
57    Path('bench_report.json').write_text(json.dumps(report,indent=2)); print(json.dumps(report,indent=2))
58if __name__=='__main__':main()