# /// script
# requires-python = ">=3.10"
# dependencies = ["unsloth", "datasets", "trl", "huggingface_hub"]
# ///
"""Stage 1 LoRA training with Unsloth on Hugging Face Jobs (1x A100 80GB).

Run (pipeline test, 10 steps):
  hf jobs uv run --flavor a100-large --timeout 45m --secrets HF_TOKEN train/hf_train.py -- --max-steps 10
Full run: no --max-steps (2 epochs). Settings follow train/config.yaml; alpha is a proposal, see STATE.md.
"""
import argparse, json, os, time

ap = argparse.ArgumentParser()
ap.add_argument("--model", default="unsloth/Qwen3.8-27B-unsloth-bnb-4bit")
ap.add_argument("--data", default="erhankeseli/abap-stage1-data")
ap.add_argument("--out", default="erhankeseli/abap-stage1-adapter-test")
ap.add_argument("--max-steps", type=int, default=-1)
ap.add_argument("--epochs", type=float, default=2.0)
ap.add_argument("--rank", type=int, default=16)
ap.add_argument("--alpha", type=int, default=16)
ap.add_argument("--lr", type=float, default=5e-5)
ap.add_argument("--max-seq-length", type=int, default=16384)
a = ap.parse_args()

from unsloth import FastLanguageModel  # noqa: E402 (import first)
from datasets import load_dataset  # noqa: E402
from trl import SFTConfig, SFTTrainer  # noqa: E402

tok_hf = os.environ["HF_TOKEN"]
model, tok = FastLanguageModel.from_pretrained(
    a.model, max_seq_length=a.max_seq_length, load_in_4bit=True, token=tok_hf)
model = FastLanguageModel.get_peft_model(
    model, r=a.rank, lora_alpha=a.alpha, lora_dropout=0.0, bias="none",
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj",
                    "in_proj_qkv", "in_proj_z", "out_proj"],
    use_gradient_checkpointing="unsloth", random_state=20261003)

ds = load_dataset(a.data, data_files={"train": "train.jsonl", "valid": "valid.jsonl"}, token=tok_hf)
cfg = SFTConfig(
    output_dir="out", per_device_train_batch_size=1, per_device_eval_batch_size=1,
    gradient_accumulation_steps=1, num_train_epochs=a.epochs, max_steps=a.max_steps,
    learning_rate=a.lr, lr_scheduler_type="cosine", warmup_steps=min(30, max(1, a.max_steps // 3)) if a.max_steps > 0 else 30,
    optim="adamw_8bit", weight_decay=0.0, logging_steps=1, eval_strategy="no", save_strategy="no",
    max_length=a.max_seq_length, dataset_text_field="text", packing=False, seed=20261003, report_to="none")
tr = SFTTrainer(model=model, processing_class=tok, train_dataset=ds["train"], eval_dataset=ds["valid"], args=cfg)

t0 = time.time()
tr.train()
train_s = time.time() - t0
steps = tr.state.global_step
ev = tr.evaluate()  # valid loss of the adapter, for the conversion check
info = {"steps": steps, "train_seconds": round(train_s, 1), "sec_per_step": round(train_s / max(steps, 1), 2),
        "eval_loss": ev.get("eval_loss"), "rank": a.rank, "alpha": a.alpha, "lr": a.lr,
        "peak_gpu_gb": round(__import__("torch").cuda.max_memory_allocated() / 2**30, 1)}
print("RESULT", json.dumps(info))
model.save_pretrained("adapter")
json.dump(info, open("adapter/job_result.json", "w"))
from huggingface_hub import HfApi  # noqa: E402
HfApi(token=tok_hf).upload_folder(folder_path="adapter", repo_id=a.out, repo_type="model")
print("PUSHED", a.out)