# /// script
# dependencies = ["torch==2.8.0", "transformers==4.56.2", "peft==0.17.1", "accelerate==1.10.1", "huggingface_hub==0.34.4", "bitsandbytes==0.47.0", "safetensors==0.6.2"]
# ///
"""Budget-bounded V3 training with paired, disjoint evaluation and reload gate."""
import os, json, random, time, re, gc, hashlib, traceback
from pathlib import Path
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, Trainer, TrainingArguments, TrainerCallback, set_seed
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training, PeftModel
from huggingface_hub import HfApi, hf_hub_download

REPO='SEE7/Astheral'; REV='37adef9b3917e4a3de0bb77f0db71089b7b82202'
OUT=Path('/tmp/astheral-v3'); OUT.mkdir(exist_ok=True)
api=HfApi(); started=time.monotonic(); set_seed(20261113); rng=random.Random(20261113)
assert torch.cuda.is_available() and 'A10G' in torch.cuda.get_device_name(0)
branch='v3-expanded-20261113'
api.create_branch(REPO,branch=branch,revision=REV,exist_ok=True)
settings=json.loads(Path(hf_hub_download(REPO,'settings.json',revision=REV)).read_text())
tok=AutoTokenizer.from_pretrained(REPO,revision=REV); tok.pad_token_id=128004
report={'status':'preparing','parent_revision':REV,'hardware':torch.cuda.get_device_name(0),'job_timeout_seconds':5400,'listed_usd_per_hour':1.0,'training_seed':20261113}

def publish():
    report['elapsed_seconds']=round(time.monotonic()-started,2)
    (OUT/'v3_report.json').write_text(json.dumps(report,indent=2))
    api.upload_folder(repo_id=REPO,folder_path=OUT,revision=branch,path_in_repo='v3-results',ignore_patterns=['merged/**'],commit_message='Record measured V3 training and evaluation artifacts')

def sample(a,b,c,kind):
    if kind==0:return f'Return only the integer: {a} + {b} - {c}.',str(a+b-c),'integer'
    if kind==1:return f'A crate has {a} red and {b} blue beads. Remove {c} red beads. Return only the total beads remaining.',str(a+b-c),'integer'
    if kind==2:return f'Return only valid JSON with keys "item" and "count": item=part_{a}, count={b}.',{'item':f'part_{a}','count':b},'json'
    if kind==3:return f'Return only one CSV value row, no header: name=Item{a}; quantity={b}; code=K{c}.',f'Item{a},{b},K{c}','exact'
    if kind==4:return f'Copy exactly, no other text: Key_{a}-B{b}_C{c}',f'Key_{a}-B{b}_C{c}','exact'
    return f'Write only a Python function named solve(x) that returns x * {b} + {c}.',f'x * {b} + {c}','code'

def grade(case,text):
    expected=case['expected']; kind=case['kind']; text=text.strip()
    try:
        if kind=='json':return json.loads(text)==expected
        if kind=='integer':return re.fullmatch(r'-?\d+',text) is not None and int(text)==int(expected)
        if kind=='exact':return text==expected
        # Parse and evaluate a tiny arithmetic-only AST; never execute model-generated code.
        import ast
        clean=re.sub(r'^```(?:python)?\s*|\s*```$','',text).strip();tree=ast.parse(clean)
        if len(tree.body)!=1 or not isinstance(tree.body[0],ast.FunctionDef):return False
        fn=tree.body[0]
        if fn.name!='solve' or len(fn.args.args)!=1 or fn.args.args[0].arg!='x' or len(fn.body)!=1 or not isinstance(fn.body[0],ast.Return):return False
        def calc(n,x):
            if isinstance(n,ast.Constant) and isinstance(n.value,(int,float)):return n.value
            if isinstance(n,ast.Name) and n.id=='x':return x
            if isinstance(n,ast.BinOp) and isinstance(n.op,(ast.Add,ast.Sub,ast.Mult)):
                a,b=calc(n.left,x),calc(n.right,x)
                return a+b if isinstance(n.op,ast.Add) else a-b if isinstance(n.op,ast.Sub) else a*b
            raise ValueError('Unsupported expression')
        match=re.fullmatch(r'x \* (\d+) \+ (\d+)',expected);b,c=map(int,match.groups())
        return all(calc(fn.body[0].value,x)==x*b+c for x in [-5,0,1,7,100])
    except Exception:return False

tests=[]
for i in range(60):
    # Evaluation numbers are outside the synthetic training ranges.
    a,b,c=701+i*3,91+i,31+i%17;p,e,k=sample(a,b,c,i%6)
    tests.append({'id':f'heldout_{i:02}','prompt':p,'expected':e,'kind':k})
train=[]
for i in range(900):
    a,b,c=rng.randint(100,500),rng.randint(10,80),rng.randint(1,29)
    p,e,k=sample(a,b,c,i%6)
    answer=json.dumps(e) if k=='json' else 'def solve(x):\n    return '+e if k=='code' else str(e)
    train.append({'messages':[{'role':'user','content':p},{'role':'assistant','content':answer}]})
retained=[json.loads(s) for s in Path(hf_hub_download(REPO,'train.jsonl',revision=REV)).read_text().splitlines()]
rng.shuffle(retained);train+=retained[:900];rng.shuffle(train)
assert not {x['prompt'] for x in tests} & {m['content'] for row in train for m in row['messages'] if m['role']=='user'}
(OUT/'heldout_cases.json').write_text(json.dumps(tests,indent=2));(OUT/'training.jsonl').write_text('\n'.join(json.dumps(x) for x in train)+'\n')
(OUT/'train_v3.py').write_text(Path(__file__).read_text())
q=BitsAndBytesConfig(load_in_4bit=True,bnb_4bit_quant_type='nf4',bnb_4bit_use_double_quant=True,bnb_4bit_compute_dtype=torch.bfloat16)
model=AutoModelForCausalLM.from_pretrained(REPO,revision=REV,quantization_config=q,device_map={'':'cuda:0'},torch_dtype=torch.bfloat16,attn_implementation='sdpa')

@torch.inference_mode()
def evaluate(label):
    model.eval(); rows=[]
    for i,case in enumerate(tests):
        inputs=tok.apply_chat_template([{'role':'system','content':settings['system']},{'role':'user','content':case['prompt']}],tokenize=True,add_generation_prompt=True,return_dict=True,return_tensors='pt').to('cuda')
        out=model.generate(**inputs,max_new_tokens=96,do_sample=False,pad_token_id=128004,eos_token_id=128009)
        text=tok.decode(out[0,inputs['input_ids'].shape[1]:],skip_special_tokens=True)
        rows.append({**case,'output':text,'pass':grade(case,text),'tokens':int(out.shape[1]-inputs['input_ids'].shape[1])})
        print('EVAL',label,i,rows[-1]['pass'],flush=True)
    (OUT/(label+'_outputs.json')).write_text(json.dumps(rows,indent=2))
    return rows

try:
    baseline=evaluate('v2_quantized'); report['baseline_passed']=sum(r['pass'] for r in baseline);publish()
    model=prepare_model_for_kbit_training(model,use_gradient_checkpointing=True)
    model=get_peft_model(model,LoraConfig(r=32,lora_alpha=64,lora_dropout=0.05,target_modules=['q_proj','k_proj','v_proj','o_proj','gate_proj','up_proj','down_proj'],task_type='CAUSAL_LM'))
    model.config.use_cache=False;encoded=[]
    for row in train:
        messages=row['messages']; prefix=tok.apply_chat_template(messages[:-1],tokenize=True,add_generation_prompt=True)
        full=tok.apply_chat_template(messages,tokenize=True,add_generation_prompt=False)
        if len(full)<=512 and full[:len(prefix)]==prefix:encoded.append({'input_ids':full,'labels':[-100]*len(prefix)+full[len(prefix):]})
    def collate(rows):
        length=max(len(r['input_ids']) for r in rows)
        return {k:torch.tensor([r[k]+[pad]*(length-len(r[k])) for r in rows]) for k,pad in [('input_ids',128004),('labels',-100)]}|{'attention_mask':torch.tensor([[1]*len(r['input_ids'])+[0]*(length-len(r['input_ids'])) for r in rows])}
    class Budget(TrainerCallback):
        def on_step_end(self,args,state,control,**kwargs):
            if time.monotonic()-started>2400:control.should_training_stop=True
            return control
    args=TrainingArguments(output_dir='/tmp/v3-checkpoints',max_steps=360,per_device_train_batch_size=1,gradient_accumulation_steps=4,learning_rate=5e-6,warmup_ratio=0.05,bf16=True,gradient_checkpointing=True,gradient_checkpointing_kwargs={'use_reentrant':False},optim='paged_adamw_8bit',save_strategy='no',logging_steps=10,report_to='none',remove_unused_columns=False,seed=20261113)
    trainer=Trainer(model=model,args=args,train_dataset=encoded,data_collator=collate,callbacks=[Budget()])
    report.update(status='training',training_examples=len(encoded),trainable_parameters=sum(p.numel() for p in model.parameters() if p.requires_grad));publish()
    result=trainer.train();report['training']=result.metrics;report['optimizer_steps']=trainer.state.global_step
    adapter=OUT/'adapter';model.save_pretrained(adapter,safe_serialization=True);tok.save_pretrained(adapter)
    model.config.use_cache=True;model.gradient_checkpointing_disable();candidate=evaluate('v3_quantized')
    report['candidate_passed']=sum(r['pass'] for r in candidate);report['status']='trained';publish()
    del trainer,model;gc.collect();torch.cuda.empty_cache()
    model=AutoModelForCausalLM.from_pretrained(REPO,revision=REV,torch_dtype=torch.bfloat16,device_map={'':'cuda:0'},attn_implementation='sdpa')
    model=PeftModel.from_pretrained(model,adapter).merge_and_unload(safe_merge=True)
    model.config.astheral_version='3';merged=OUT/'merged';model.save_pretrained(merged,safe_serialization=True,max_shard_size='4GB');tok.save_pretrained(merged)
    del model;gc.collect();torch.cuda.empty_cache()
    model=AutoModelForCausalLM.from_pretrained(merged,torch_dtype=torch.bfloat16,device_map={'':'cuda:0'},attn_implementation='sdpa')
    reloaded=evaluate('v3_reloaded');report['reloaded_passed']=sum(r['pass'] for r in reloaded)
    groups={kind:{'baseline':sum(r['pass'] for r in baseline if r['kind']==kind),'candidate':sum(r['pass'] for r in reloaded if r['kind']==kind),'total':sum(r['kind']==kind for r in tests)} for kind in {x['kind'] for x in tests}}
    report['groups']=groups;report['quality_gate_passed']=report['reloaded_passed']>report['baseline_passed'] and all(v['candidate']>=v['baseline'] for v in groups.values())
    report['status']='complete';report['peak_gpu_bytes']=torch.cuda.max_memory_allocated();publish()
    # Save the actual trained weights to the isolated candidate branch even on a failed gate.
    api.upload_folder(repo_id=REPO,folder_path=merged,revision=branch,commit_message='Save reloaded V3 candidate weights; see measured quality gate in v3-results')
    if report['quality_gate_passed']:
        api.create_tag(REPO,tag='v3-expanded-candidate',revision=api.model_info(REPO,revision=branch).sha,exist_ok=True)
    print('V3_COMPLETE',json.dumps(report),flush=True)
except Exception:
    report['status']='failed';report['error']=traceback.format_exc();publish();raise