Upload opi_mathqa_full.py

opi_mathqa_full.py

@@ -0,0 +1,268 @@
+#!/usr/bin/env python3
+"""Full OPi-vs-GELU Math-QA experiment on math-ai/StackMathQA.
+
+Design:
+  * StackMathQA 100K/200K/400K/800K/1600K
+  * question-level leakage-safe split
+  * SentencePiece BPE tokenizer trained on TRAIN ONLY
+  * identical encoder-decoder Transformer baseline vs OPi FFN
+  * OPi activation derived from k(x)=log|cos(pi*x/log(x))|
+    with a smooth squared-cosine form and a sign-preserving multiplicative gate
+  * teacher-forced validation loss/perplexity
+  * greedy generation on held-out questions
+  * token accuracy, normalized exact match, character F1, ROUGE-L
+  * numerical-answer extraction/equivalence when possible
+  * checkpointing, JSON report, CSV predictions, SHA-256/SHA-512
+
+The full 100K file is 181 MB. StackMathQA also provides 200K, 400K, 800K and
+1.6M configurations. The official dataset describes Q and A as LaTeX-encoded
+math text and CC-BY-4.0.
+"""
+import argparse, csv, hashlib, json, math, os, random, re, time
+from collections import Counter
+from pathlib import Path
+from urllib.request import Request, urlopen
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import sentencepiece as spm
+
+SEED=42
+random.seed(SEED); torch.manual_seed(SEED)
+
+CONFIGS={
+ '100k':'stackmathqa100k', '200k':'stackmathqa200k', '400k':'stackmathqa400k',
+ '800k':'stackmathqa800k', '1600k':'stackmathqa1600k'}
+BASE='https://huggingface.co/datasets/math-ai/StackMathQA/resolve/main/data/{cfg}/all.jsonl'
+EXPECTED={
+ '100k':'ddf7526cb21600dcf93991e400393eb9adc3549a6077d0cd0b0d37c45e1b037a'}
+
+SPECIAL=['<pad>','<bos>','<eos>','<unk>']
+
+
+def set_seed(seed):
+ random.seed(seed); torch.manual_seed(seed); torch.cuda.manual_seed_all(seed)
+
+
+def sha(path, alg='sha256'):
+ h=hashlib.new(alg)
+ with open(path,'rb') as f:
+  for b in iter(lambda:f.read(1024*1024),b''): h.update(b)
+ return h.hexdigest()
+
+
+def download(url,path):
+ path=Path(path); path.parent.mkdir(parents=True,exist_ok=True)
+ if path.exists(): return
+ req=Request(url,headers={'User-Agent':'OPi-MathQA/2.0'})
+ with urlopen(req,timeout=60) as r, open(path,'wb') as f:
+  total=int(r.headers.get('Content-Length') or 0); done=0; t=time.time()
+  while True:
+   b=r.read(4*1024*1024)
+   if not b: break
+   f.write(b); done+=len(b)
+   if total and time.time()-t>2:
+    print(f'\rdownload {done/1e6:.1f}/{total/1e6:.1f} MB',end='',flush=True); t=time.time()
+ print()
+
+
+def norm(s):
+ s=s.replace('\r',' ').replace('\n',' ')
+ s=re.sub(r'\s+',' ',s).strip().lower()
+ return s
+
+
+def canonical_question(q):
+ # Remove harmless whitespace differences while preserving LaTeX symbols.
+ return norm(q)
+
+
+def load_rows(path,limit=0):
+ out=[]; seen=set(); questions=set()
+ with open(path,encoding='utf-8') as f:
+  for line in f:
+   if not line.strip(): continue
+   x=json.loads(line); q=(x.get('Q') or '').strip(); a=(x.get('A') or '').strip()
+   if not q or not a: continue
+   qk=canonical_question(q); k=(qk,norm(a))
+   if k in seen: continue
+   seen.add(k); out.append((q,a,x.get('meta',{})))
+   questions.add(qk)
+   if limit and len(out)>=limit: break
+ return out
+
+
+def split(rows, seed=SEED, test=.10, val=.10):
+ # All answers for the same question stay in one split.
+ out={'train':[],'val':[],'test':[]}
+ for row in rows:
+  qk=canonical_question(row[0])
+  h=int(hashlib.sha256((str(seed)+'|'+qk).encode()).hexdigest()[:12],16)/float(16**12-1)
+  out['test' if h<test else 'val' if h<test+val else 'train'].append(row)
+ return out
+
+
+def train_spm(train_rows, path, vocab_size=16000):
+ path=Path(path); path.parent.mkdir(parents=True,exist_ok=True)
+ if path.exists(): return
+ corpus=path.with_suffix('.txt')
+ with open(corpus,'w',encoding='utf-8') as f:
+  for q,a,_ in train_rows:
+   f.write(q.replace('\n',' ')+'\n'); f.write(a.replace('\n',' ')+'\n')
+ spm.SentencePieceTrainer.Train(
+  input=str(corpus), model_prefix=str(path.with_suffix('')),
+  vocab_size=vocab_size, model_type='bpe', character_coverage=1.0,
+  pad_id=0, bos_id=1, eos_id=2, unk_id=3,
+  user_defined_symbols=[], hard_vocab_limit=False)
+ corpus.unlink(missing_ok=True)
+
+
+class Tok:
+ def __init__(self, model):
+  self.sp=spm.SentencePieceProcessor(model_file=str(model))
+  self.pad=self.sp.pad_id(); self.bos=self.sp.bos_id(); self.eos=self.sp.eos_id(); self.unk=self.sp.unk_id()
+ def enc(self,s,maxlen):
+  ids=[self.bos]+self.sp.encode(s,out_type=int)[:maxlen-2]+[self.eos]
+  return ids+[self.pad]*(maxlen-len(ids))
+ def dec(self,ids):
+  ids=[i for i in ids if i not in (self.pad,self.bos)]
+  if self.eos in ids: ids=ids[:ids.index(self.eos)]
+  return self.sp.decode(ids)
+ def __len__(self): return self.sp.get_piece_size()
+
+
+def batch(rows,tok,maxq,maxa,device):
+ q=torch.tensor([tok.enc(r[0],maxq) for r in rows],dtype=torch.long,device=device)
+ a=torch.tensor([tok.enc(r[1],maxa) for r in rows],dtype=torch.long,device=device)
+ return q,a
+
+
+class OPiActivation(nn.Module):
+ """Sign-preserving activation derived from OPi's phase/kernel.
+
+ k(x)=log|cos(u(x))|, u(x)=pi*x/log(x).
+ We use 0.5*log(cos(u)^2+eps), a smooth numerical version, then use the
+ kernel as a multiplicative gate on z so the activation retains sign/information.
+ """
+ def __init__(self, dim, alpha=.65, beta=.75, offset=1.5, eps=1e-5):
+  super().__init__(); self.register_buffer('alpha',torch.tensor(float(alpha))); self.register_buffer('beta',torch.tensor(float(beta))); self.offset=offset; self.eps=eps
+ def forward(self,z):
+  x=self.offset+F.softplus(z)
+  u=math.pi*x/torch.log(x)
+  k=.5*torch.log(torch.cos(u).square()+self.eps)
+  gate=1.0 + torch.clamp(self.beta, -0.95, 0.95)*torch.tanh(self.alpha*k)
+  return z*gate
+
+
+class FFN(nn.Module):
+ def __init__(self,d,ff,opi):
+  super().__init__(); self.fc1=nn.Linear(d,ff); self.act=OPiActivation(ff) if opi else nn.GELU(); self.fc2=nn.Linear(ff,d)
+ def forward(self,x): return self.fc2(self.act(self.fc1(x)))
+
+class EncBlock(nn.Module):
+ def __init__(self,d,heads,ff,drop,opi):
+  super().__init__(); self.att=nn.MultiheadAttention(d,heads,batch_first=True,dropout=drop); self.n1=nn.LayerNorm(d); self.ff=FFN(d,ff,opi); self.n2=nn.LayerNorm(d); self.drop=nn.Dropout(drop)
+ def forward(self,x,padmask):
+  y,_=self.att(x,x,x,key_padding_mask=padmask,need_weights=False); x=self.n1(x+self.drop(y)); return self.n2(x+self.drop(self.ff(x)))
+
+class DecBlock(nn.Module):
+ def __init__(self,d,heads,ff,drop,opi):
+  super().__init__(); self.selfatt=nn.MultiheadAttention(d,heads,batch_first=True,dropout=drop); self.cross=nn.MultiheadAttention(d,heads,batch_first=True,dropout=drop); self.n1=nn.LayerNorm(d); self.n2=nn.LayerNorm(d); self.n3=nn.LayerNorm(d); self.ff=FFN(d,ff,opi); self.drop=nn.Dropout(drop)
+ def forward(self,x,mem,srcmask,tgtmask):
+  y,_=self.selfatt(x,x,x,attn_mask=tgtmask,need_weights=False); x=self.n1(x+self.drop(y))
+  y,_=self.cross(x,mem,mem,key_padding_mask=srcmask,need_weights=False); x=self.n2(x+self.drop(y)); return self.n3(x+self.drop(self.ff(x)))
+
+class Model(nn.Module):
+ def __init__(self,vocab,d=256,heads=8,ff=512,layers=4,drop=.1,opi=False,maxlen=512):
+  super().__init__(); self.d=d; self.maxlen=maxlen; self.emb=nn.Embedding(vocab,d,padding_idx=0); self.pos=nn.Embedding(maxlen,d); self.enc=nn.ModuleList([EncBlock(d,heads,ff,drop,opi) for _ in range(layers)]); self.dec=nn.ModuleList([DecBlock(d,heads,ff,drop,opi) for _ in range(layers)]); self.out=nn.Linear(d,vocab)
+ def forward(self,src,tgt):
+  S=src.size(1); T=tgt.size(1); ps=torch.arange(S,device=src.device)[None,:]; pt=torch.arange(T,device=tgt.device)[None,:]
+  mem=self.emb(src)+self.pos(ps); srcmask=src.eq(0)
+  for b in self.enc: mem=b(mem,srcmask)
+  y=self.emb(tgt)+self.pos(pt); mask=torch.triu(torch.ones(T,T,device=tgt.device,dtype=torch.bool),1)
+  for b in self.dec: y=b(y,mem,srcmask,mask)
+  return self.out(y)
+
+
+def lcs(a,b):
+ prev=[0]*(len(b)+1)
+ for x in a:
+  cur=[0]
+  for j,y in enumerate(b,1): cur.append(prev[j-1]+1 if x==y else max(prev[j],cur[-1]))
+  prev=cur
+ return prev[-1]
+
+
+def charf(ref,hyp):
+ rc=Counter(ref); hc=Counter(hyp); ov=sum((rc&hc).values())
+ if not ref or not hyp: return 0.0
+ p=ov/len(hyp); r=ov/len(ref); return 2*p*r/(p+r) if p+r else 0.0
+
+NUM_RE=re.compile(r'(?<![A-Za-z])[-+]?\d+(?:\.\d+)?(?:e[-+]?\d+)?',re.I)
+
+def nums(s):
+ return [round(float(x),10) for x in NUM_RE.findall(s)]
+
+def metric_rows(refs,hyps):
+ tok_ok=tok_total=0; exact=0; cf=[]; rl=[]; numeq=[]
+ for r,h in zip(refs,hyps):
+  rt=norm(r).split(); ht=norm(h).split(); n=min(len(rt),len(ht)); tok_ok+=sum(rt[i]==ht[i] for i in range(n)); tok_total+=len(rt)
+  exact+=norm(r)==norm(h); cf.append(charf(norm(r),norm(h))); L=lcs(rt,ht); rl.append(L/max(1,len(rt)))
+  rn,hn=nums(r),nums(h); numeq.append(bool(rn) and rn==hn)
+ return {'token_accuracy':tok_ok/max(1,tok_total),'exact_match':exact/max(1,len(refs)),'char_f1':sum(cf)/max(1,len(cf)),'rougeL_recall':sum(rl)/max(1,len(rl)),'numeric_answer_match_when_both_numeric':sum(numeq)/max(1,len(numeq))}
+
+
+def generate(model,rows,tok,device,maxq,maxa,bs=4):
+ model.eval(); refs=[]; hyps=[]
+ with torch.no_grad():
+  for i in range(0,len(rows),bs):
+   chunk=rows[i:i+bs]; src=torch.tensor([tok.enc(r[0],maxq) for r in chunk],dtype=torch.long,device=device); B=src.size(0); out=torch.full((B,1),tok.bos,dtype=torch.long,device=device)
+   for _ in range(maxa-1):
+    logits=model(src,out); nxt=logits[:,-1].argmax(-1); out=torch.cat([out,nxt[:,None]],1)
+    if bool((nxt==tok.eos).all()): break
+   for j,r in enumerate(chunk): refs.append(r[1]); hyps.append(tok.dec(out[j].tolist()))
+ return refs,hyps
+
+
+def eval_loss(model,rows,tok,device,maxq,maxa,bs):
+ model.eval(); lossf=nn.CrossEntropyLoss(ignore_index=tok.pad); total=n=0
+ with torch.no_grad():
+  for i in range(0,len(rows),bs):
+   q,a=batch(rows[i:i+bs],tok,maxq,maxa,device); z=model(q,a[:,:-1]); loss=lossf(z.reshape(-1,z.size(-1)),a[:,1:].reshape(-1)); total+=float(loss)*q.size(0); n+=q.size(0)
+ return total/max(1,n)
+
+
+def train(model,train_rows,val_rows,tok,args,device):
+ opt=torch.optim.AdamW(model.parameters(),lr=args.lr,weight_decay=.01); lossf=nn.CrossEntropyLoss(ignore_index=tok.pad); best=float('inf'); best_state=None; history=[]
+ for ep in range(1,args.epochs+1):
+  model.train(); random.shuffle(train_rows); total=n=0; t=time.time()
+  for i in range(0,len(train_rows),args.batch):
+   q,a=batch(train_rows[i:i+args.batch],tok,args.max_q,args.max_a,device); z=model(q,a[:,:-1]); loss=lossf(z.reshape(-1,z.size(-1)),a[:,1:].reshape(-1)); opt.zero_grad(set_to_none=True); loss.backward(); nn.utils.clip_grad_norm_(model.parameters(),1.0); opt.step(); total+=float(loss)*q.size(0); n+=q.size(0)
+  vl=eval_loss(model,val_rows[:args.val_eval],tok,device,args.max_q,args.max_a,args.eval_batch); rec={'epoch':ep,'train_loss':total/max(1,n),'val_loss':vl,'val_ppl':math.exp(min(vl,20)),'seconds':time.time()-t}; history.append(rec); print(json.dumps(rec),flush=True)
+  if vl<best: best=vl; best_state={k:v.detach().cpu().clone() for k,v in model.state_dict().items()};
+  if args.patience and ep-history.index(min(history,key=lambda x:x['val_loss']))+1>=args.patience and ep>args.patience: break
+ if best_state: model.load_state_dict(best_state)
+ return history
+
+
+def main():
+ ap=argparse.ArgumentParser(); ap.add_argument('--config',choices=CONFIGS,default='100k'); ap.add_argument('--data',default=''); ap.add_argument('--download',action='store_true'); ap.add_argument('--rows',type=int,default=0); ap.add_argument('--epochs',type=int,default=3); ap.add_argument('--batch',type=int,default=16); ap.add_argument('--eval-batch',type=int,default=8); ap.add_argument('--val-eval',type=int,default=2000); ap.add_argument('--test-eval',type=int,default=2000); ap.add_argument('--max-q',type=int,default=384); ap.add_argument('--max-a',type=int,default=512); ap.add_argument('--vocab',type=int,default=16000); ap.add_argument('--d',type=int,default=256); ap.add_argument('--heads',type=int,default=8); ap.add_argument('--ff',type=int,default=512); ap.add_argument('--layers',type=int,default=4); ap.add_argument('--lr',type=float,default=2e-4); ap.add_argument('--patience',type=int,default=2); ap.add_argument('--device',default='cpu'); ap.add_argument('--out',default='opi_mathqa_report.json'); ap.add_argument('--predictions',default='opi_mathqa_predictions.csv'); ap.add_argument('--workdir',default='opi_mathqa_work'); ap.add_argument('--seed',type=int,default=42)
+ args=ap.parse_args(); set_seed(args.seed)
+ work=Path(args.workdir); work.mkdir(parents=True,exist_ok=True); path=Path(args.data or work/f'{CONFIGS[args.config]}.jsonl')
+ if args.download: download(BASE.format(cfg=CONFIGS[args.config]),path)
+ if not path.exists(): raise SystemExit(f'Dataset missing: {path}. Use --download on a network-enabled machine.')
+ got=sha(path); print('dataset_sha256',got); 
+ if args.config in EXPECTED and not args.data: print('expected_sha256',EXPECTED[args.config],'match=',got==EXPECTED[args.config])
+ rows=load_rows(path,args.rows); splits=split(rows,args.seed); print('rows',len(rows),'splits',{k:len(v) for k,v in splits.items()})
+ spm_path=work/'tokenizer.model'; train_spm(splits['train'],spm_path,args.vocab); tok=Tok(spm_path); print('vocab',len(tok))
+ device=torch.device(args.device); report={'dataset':{'name':'math-ai/StackMathQA','config':args.config,'rows':len(rows),'splits':{k:len(v) for k,v in splits.items()},'sha256':got,'sha512':sha(path,'sha512')},'model':vars(args),'results':{}}
+ for name,opi in [('baseline_gelu',False),('opi',True)]:
+  set_seed(args.seed); model=Model(len(tok),args.d,args.heads,args.ff,args.layers,.1,opi,max(args.max_q,args.max_a)).to(device); params=sum(p.numel() for p in model.parameters()); print(name,'params',params)
+  hist=train(model,splits['train'],splits['val'],tok,args,device); refs,hyps=generate(model,splits['test'][:args.test_eval],tok,device,args.max_q,args.max_a,args.eval_batch); m=metric_rows(refs,hyps); m['test_loss']=eval_loss(model,splits['test'][:args.test_eval],tok,device,args.max_q,args.max_a,args.eval_batch); m['ppl']=math.exp(min(m['test_loss'],20)); m['params']=params; m['history']=hist; report['results'][name]=m
+  with open(args.predictions.replace('.csv',f'_{name}.csv'),'w',newline='',encoding='utf-8') as f:
+   w=csv.writer(f); w.writerow(['question','reference','prediction']); w.writerows(zip([r[0] for r in splits['test'][:args.test_eval]],refs,hyps))
+  torch.save({'state_dict':model.state_dict(),'tokenizer':str(spm_path),'args':vars(args),'metrics':m},work/f'{name}.pt')
+ Path(args.out).write_text(json.dumps(report,ensure_ascii=False,indent=2),encoding='utf-8'); print(json.dumps(report['results'],indent=2))
+
+if __name__=='__main__': main()