"""Astheral v2: remote L4 QLoRA training and measured release."""
import os,json,time,gc,hashlib,shutil,math,random,urllib.request,collections
from pathlib import Path
import torch
from torch.utils.data import Dataset
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
from evaluate import CASES,check,summary
ROOT=Path(__file__).resolve().parent
OUT=Path("/tmp/astheral-v2-output");OUT.mkdir(exist_ok=True)
REPO="SEE7/Astheral";BASE="NousResearch/DeepHermes-3-Llama-3-8B-Preview";REV="53d902c7dcfa4c749a3b455149064f82af2e2549"
api=HfApi();start=time.monotonic();set_seed(20261113)
assert torch.cuda.is_available(),"Remote CUDA GPU required"
assert "L4" in torch.cuda.get_device_name(0),"This job is provisioned for remote L4"
print("REMOTE_HARDWARE",torch.cuda.get_device_name(0),flush=True)
settings=json.loads((ROOT/"settings.json").read_text())
tok=AutoTokenizer.from_pretrained(BASE,revision=REV);tok.pad_token_id=128004
def messages(prompt,mode):
    return [{"role":"system","content":settings["system"]+("\n"+settings["thinking"] if mode=="thinking" else "")},{"role":"user","content":prompt}]
class Samples(Dataset):
    def __init__(self,rows):
        self.items=[];self.kept=[];self.skipped=0
        evalprompts={c["prompt"] for c in CASES}
        for r in rows:
            if any(m["role"]=="user" and m["content"] in evalprompts for m in r["messages"]):raise RuntimeError("Evaluation overlap")
            prompt=tok.apply_chat_template(r["messages"][:-1],tokenize=True,add_generation_prompt=True)
            full=tok.apply_chat_template(r["messages"],tokenize=True,add_generation_prompt=False)
            if len(full)>1280:self.skipped+=1;continue
            assert full[:len(prompt)]==prompt
            self.items.append({"input_ids":full,"labels":[-100]*len(prompt)+full[len(prompt):]})
            self.kept.append(r)
    def __len__(self):return len(self.items)
    def __getitem__(self,i):return self.items[i]
def collate(batch):
    size=max(len(x["input_ids"]) for x in batch)
    return {k:torch.tensor([x[k]+[fill]*(size-len(x[k])) for x in batch]) for k,fill in [("input_ids",128004),("labels",-100)]}|{"attention_mask":torch.tensor([[1]*len(x["input_ids"])+[0]*(size-len(x["input_ids"])) for x in batch])}
records=[json.loads(s) for s in (ROOT/"curriculum.jsonl").read_text().splitlines()]
random.Random(20261113).shuffle(records)
# Loss-validation conversations are disjoint from training.
valid=Samples(records[:64]);train=Samples(records[64:])
assert len(train)>1800
(OUT/"train.jsonl").write_text("\n".join(json.dumps(r) for r in train.kept)+"\n")
(OUT/"validation.jsonl").write_text("\n".join(json.dumps(r) for r in valid.kept)+"\n")
print("DATA_READY",len(train),len(valid),"skipped_long",train.skipped,flush=True)
quant=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(BASE,revision=REV,torch_dtype=torch.bfloat16,device_map={"":"cuda:0"},attn_implementation="sdpa")
@torch.inference_mode()
def generate(tag,cases=CASES,batch_size=4):
    model.eval();tok.padding_side="left";results=[]
    for mode in ["instruct","thinking"]:
        selected=[c for c in cases if c["mode"]==mode]
        for i in range(0,len(selected),batch_size):
            group=selected[i:i+batch_size];set_seed(667+i)
            texts=[tok.apply_chat_template(messages(c["prompt"],mode),tokenize=False,add_generation_prompt=True) for c in group]
            inputs=tok(texts,padding=True,return_tensors="pt").to("cuda")
            out=model.generate(**inputs,max_new_tokens=1024 if mode=="thinking" else 384,do_sample=False,pad_token_id=128004,eos_token_id=128009)
            for c,o in zip(group,out):
                answer=tok.decode(o[inputs["input_ids"].shape[1]:],skip_special_tokens=True)
                results.append({**c,"output":answer,"pass":check(c,answer)})
                print("EVAL",tag,c["id"],results[-1]["pass"],flush=True)
    tok.padding_side="right"
    (OUT/(tag+"_outputs.json")).write_text(json.dumps(results,indent=2))
    print("SCORE",tag,json.dumps(summary(results)),flush=True)
    return results
baseline=generate("base",batch_size=2)
del model;gc.collect();torch.cuda.empty_cache()
model=AutoModelForCausalLM.from_pretrained(BASE,revision=REV,quantization_config=quant,torch_dtype=torch.bfloat16,device_map={"":"cuda:0"},attn_implementation="sdpa")
model=prepare_model_for_kbit_training(model,use_gradient_checkpointing=True,gradient_checkpointing_kwargs={"use_reentrant":False})
model=get_peft_model(model,LoraConfig(r=32,lora_alpha=64,lora_dropout=0.05,target_modules="all-linear",task_type="CAUSAL_LM",bias="none"))
model.config.use_cache=False
trainable=sum(p.numel() for p in model.parameters() if p.requires_grad)
model.print_trainable_parameters()
class Budget(TrainerCallback):
    def on_step_end(self,args,state,control,**kwargs):
        if time.monotonic()-start>7200:control.should_training_stop=True
        return control
args=TrainingArguments(output_dir="/tmp/astheral-v2-checkpoints",num_train_epochs=1.5,per_device_train_batch_size=1,per_device_eval_batch_size=1,gradient_accumulation_steps=8,learning_rate=4e-5,lr_scheduler_type="cosine",warmup_ratio=0.05,weight_decay=0.01,bf16=True,gradient_checkpointing=True,gradient_checkpointing_kwargs={"use_reentrant":False},optim="paged_adamw_8bit",logging_steps=10,save_strategy="no",report_to="none",remove_unused_columns=False,seed=20261113,data_seed=20261113,max_grad_norm=0.5,dataloader_num_workers=0)
trainer=Trainer(model=model,args=args,train_dataset=train,eval_dataset=valid,data_collator=collate,callbacks=[Budget()])
before=trainer.evaluate();result=trainer.train();after=trainer.evaluate()
model.config.use_cache=True;model.gradient_checkpointing_disable()
candidate=generate("astheral")
def gate(outputs):
    s=summary(outputs)
    return s["regression"]["passed"]==s["regression"]["total"] and s["fresh"]["passed"]>=max(31,summary(baseline)["fresh"]["passed"]) and s["number"]["passed"]>=14 and s["code"]["passed"]>=5 and s["exact"]["passed"]+s["json"]["passed"]>=12
quality=gate(candidate)
# Save the adapter and reports before any merge, regardless of evaluation result.
adapter=OUT/"adapter";model.save_pretrained(adapter,safe_serialization=True);tok.save_pretrained(adapter)
metrics={"version":"2","base_model":BASE,"base_revision":REV,"training_examples":len(train),"validation_examples":len(valid),"skipped_long_examples":train.skipped,"trainable_parameters":trainable,"rank":32,"target_modules":"all-linear","quantization":"NF4 QLoRA; BF16 adapter computation","optimizer_steps":trainer.state.global_step,"epochs":trainer.state.epoch,"learning_rate":4e-5,"validation_loss_before":before["eval_loss"],"validation_loss_after":after["eval_loss"],"training":result.metrics,"base_scores":summary(baseline),"astheral_scores":summary(candidate),"quality_gate_passed":quality,"hardware":torch.cuda.get_device_name(0),"max_gpu_memory_allocated_bytes":torch.cuda.max_memory_allocated(),"elapsed_seconds":time.monotonic()-start,"generation":{"sampling":False,"instruct_max_new_tokens":384,"thinking_max_new_tokens":1024},"evaluation_design":"36 frozen fresh checks, two known regression cases, six review-only samples. Same settings for base and fine-tune."}
(OUT/"evaluation.json").write_text(json.dumps(metrics,indent=2))
(OUT/"trainer_log.json").write_text(json.dumps(trainer.state.log_history,indent=2))
for name in ["train.py","prepare_data.py","evaluate.py","settings.json","requirements.txt","data_sources.json"]:
    shutil.copy2(ROOT/name,OUT/name)
api.upload_folder(repo_id=REPO,folder_path=OUT,path_in_repo="v2-results",commit_message="Record v2 trained adapter and complete frozen-suite comparison")
print("QUALITY_GATE",quality,json.dumps(metrics),flush=True)
if not quality:
    print("QUANTIZED_GATE_MISSED; checking the final BF16 merge before release decision",flush=True)
# Merge into original BF16 weights, rather than dequantized training weights.
del trainer,model;gc.collect();torch.cuda.empty_cache()
base=AutoModelForCausalLM.from_pretrained(BASE,revision=REV,torch_dtype=torch.bfloat16,device_map={"":"cuda:0"},attn_implementation="sdpa")
model=PeftModel.from_pretrained(base,adapter).merge_and_unload(safe_merge=True)
model.config._name_or_path="Llama-3.1-Astheral"
model.config.astheral_base_model=BASE;model.config.astheral_version="2";model.config.astheral_modes=["instruct","thinking"]
# Check the known failures again after conversion back to BF16.
merged_checks=generate("merged_regressions",[c for c in CASES if c["group"]=="regression"],batch_size=1)
if not all(c["pass"] for c in merged_checks):raise RuntimeError("Merged checkpoint regression check failed")
model.save_pretrained(OUT,safe_serialization=True,max_shard_size="4GB");tok.save_pretrained(OUT)
del model,base;gc.collect();torch.cuda.empty_cache()
model=AutoModelForCausalLM.from_pretrained(OUT,torch_dtype=torch.bfloat16,device_map={"":"cuda:0"},attn_implementation="sdpa")
merged_outputs=generate("astheral_merged",batch_size=2)
if not gate(merged_outputs):raise RuntimeError("Reloaded merged model did not pass release gate")
metrics["merged_scores"]=summary(merged_outputs)
metrics["quality_gate_passed"]=True
metrics["training_source_counts"]=dict(collections.Counter(r["source"] for r in train.kept))
metrics["elapsed_seconds"]=time.monotonic()-start
(OUT/"evaluation.json").write_text(json.dumps(metrics,indent=2))
(OUT/"reload_smoke_test.json").write_text(json.dumps({"loaded_from":"saved merged BF16 checkpoint","regressions_passed":True,"full_suite_gate_passed":True},indent=2))
(OUT/"DATA.md").write_text("# Training data\n\nExact selected examples are in train.jsonl and validation.jsonl. Source revisions, component counts and dataset licenses are recorded in data_sources.json. Original Astheral examples are CC0-1.0. Model weights retain the applicable Llama license.\n")
for name in ["inference.py","LICENSE","UPSTREAM_LLAMA3_LICENSE","NOTICE"]:
    shutil.copy2(ROOT/name if (ROOT/name).exists() else Path(hf_hub_download(REPO,name)),OUT/name)
license_urls={"DATA_LICENSE_GSM8K.txt":"https://raw.githubusercontent.com/openai/grade-school-math/master/LICENSE","DATA_LICENSE_APACHE2.txt":"https://www.apache.org/licenses/LICENSE-2.0.txt"}
for name,url in license_urls.items():
    (OUT/name).write_bytes((ROOT/name).read_bytes() if (ROOT/name).exists() else urllib.request.urlopen(url,timeout=60).read())
(OUT/"DATA_LICENSE_ODCBY.txt").write_text("Contains information from bigcode/self-oss-instruct-sc2-exec-filter-50k, made available under the ODC Attribution License v1.0: https://opendatacommons.org/licenses/by/1-0/. Source: https://huggingface.co/datasets/bigcode/self-oss-instruct-sc2-exec-filter-50k\n")
(OUT/"NOTICE").write_text((OUT/"NOTICE").read_text()+"\nContains data from HuggingFaceTB/SmolTalk (Apache 2.0), openai/gsm8k (MIT; Copyright (c) 2021 OpenAI), and bigcode/self-oss-instruct-sc2-exec-filter-50k, made available under the ODC Attribution License: https://opendatacommons.org/licenses/by/1-0/. Source revisions and transformations are recorded in data_sources.json and prepare_data.py.\n")
card=(ROOT/"MODEL_CARD.md").read_text()
card+="\n## Measured evaluation\n\n"+json.dumps({"base":metrics["base_scores"],"astheral":metrics["merged_scores"]},indent=2)+"\n\nFull prompts, outputs, code tests and training records are included. The 36 fresh checks cover constrained output, multi-step numerical problems and executable Python functions. The two prior failures are tracked separately as regression checks.\n"
(OUT/"README.md").write_text(card)
manifest={}
for f in OUT.rglob("*"):
    if f.is_file():
        h=hashlib.sha256()
        with f.open("rb") as stream:
            for chunk in iter(lambda:stream.read(8388608),b""):h.update(chunk)
        manifest[str(f.relative_to(OUT))]={"bytes":f.stat().st_size,"sha256":h.hexdigest()}
(OUT/"artifact_manifest.json").write_text(json.dumps(manifest,indent=2))
# Freeze historical v1 results before updating the release.
api.create_tag(REPO,tag="v1",revision=api.model_info(REPO).sha,exist_ok=True)
api.upload_folder(repo_id=REPO,folder_path=OUT,delete_patterns=["training/*","coding_smoke_test.json"],commit_message="Release Astheral v2 with balanced training, expanded adapters and passing quality gates")
api.create_tag(REPO,tag="v2",revision=api.model_info(REPO).sha,exist_ok=True)
print("PUBLISHED_V2",api.model_info(REPO).sha,flush=True)