Upload scripts/qlora_qwythos_job.py with huggingface_hub

scripts/qlora_qwythos_job.py

@@ -0,0 +1,179 @@
+# -*- coding: utf-8 -*-
+"""
+QLoRA fine-tune Qwen3.5-9B base on Solidity security data via Unsloth.
+Targets: severity calibration + tool calling (Qwythos-9B base).
+Trains on A100 via HF Jobs.
+"""
+
+import sys, os, subprocess
+
+sys.stdout.reconfigure(encoding="utf-8", errors="replace")
+sys.stderr.reconfigure(encoding="utf-8", errors="replace")
+os.environ.setdefault("PYTHONIOENCODING", "utf-8")
+
+WORKDIR = "/workspace"
+os.makedirs(WORKDIR, exist_ok=True)
+
+
+def log(msg):
+    print(msg)
+
+
+def run_cmd(cmd, timeout=None):
+    r = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
+    return r
+
+
+def main():
+    log("=" * 50)
+    log("QLoRA Fine-Tune: Qwen3.5-9B (Unsloth)")
+    log("Dataset: qwythos-sec-training-data")
+    log("Target: severity calibration + tool calling")
+    log("=" * 50)
+
+    # ================================================================
+    # Step 1/4: Install Unsloth
+    # ================================================================
+    log("\n[1/4] Installing Unsloth ...")
+    r1 = run_cmd(["pip", "install", "--quiet", "--no-cache-dir", "unsloth"], timeout=600)
+    log(f"  pip install unsloth exit={r1.returncode}")
+    if r1.returncode != 0:
+        log(f"  stderr (last 1000): {r1.stderr[-1000:]}")
+        sys.exit(1)
+    log("  [OK] Unsloth installed")
+
+    # hf_transfer for faster model upload
+    r2 = run_cmd(["pip", "install", "--quiet", "--no-cache-dir", "hf_transfer"], timeout=60)
+    log(f"  hf_transfer exit={r2.returncode}")
+
+    # ================================================================
+    # Step 2/4: Load dataset
+    # ================================================================
+    log("\n[2/4] Loading security training dataset ...")
+    from datasets import load_dataset, Dataset
+
+    ds = load_dataset("mxguru1/qwythos-sec-training-data", split="train")
+    val_ds = load_dataset("mxguru1/qwythos-sec-training-data", split="validation")
+    log(f"  train: {len(ds)} rows, val: {len(val_ds)} rows")
+
+    # Format as chat templates for Unsloth SFT
+    def format_prompt(row):
+        text = (
+            "<|im_start|>user\n" + row["prompt"] + "<|im_end|>\n"
+            "<|im_start|>assistant\n" + row["completion"] + "<|im_end|>"
+        )
+        return {"text": text}
+
+    train_ds = ds.map(format_prompt, remove_columns=ds.column_names)
+    val_ds_out = val_ds.map(format_prompt, remove_columns=val_ds.column_names)
+    log(f"  formatted {len(train_ds)} train / {len(val_ds_out)} val samples")
+
+    # ================================================================
+    # Step 3/4: Train with Unsloth
+    # ================================================================
+    log("\n[3/4] Loading Qwen3.5-9B + tokenizer (Unsloth 4-bit) ...")
+    from unsloth import FastLanguageModel
+
+    model, tokenizer = FastLanguageModel.from_pretrained(
+        model_name="Qwen/Qwen3.5-9B",
+        max_seq_length=2048,
+        load_in_4bit=True,
+        load_in_8bit=False,
+        fast_inference=False,
+        token=os.environ.get("HF_TOKEN", ""),
+    )
+    log("  model loaded (4-bit QLoRA)")
+
+    # Add LoRA adapters - all linear modules for full coverage
+    model = FastLanguageModel.get_peft_model(
+        model,
+        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",
+            "embed_tokens", "lm_head",
+        ],
+        bias="none",
+        use_gradient_checkpointing="unsloth",
+    )
+    log("  LoRA adapters attached (r=32, all linear modules)")
+
+    log("  Starting training ...")
+    from unsloth import is_bf16_supported
+    from trl import SFTTrainer
+    from transformers import TrainingArguments, DataCollatorForSeq2Seq
+
+    trainer = SFTTrainer(
+        model=model,
+        tokenizer=tokenizer,
+        train_dataset=train_ds,
+        eval_dataset=val_ds_out,
+        dataset_text_field="text",
+        max_seq_length=2048,
+        data_collator=DataCollatorForSeq2Seq(tokenizer, model=model, padding=True),
+        args=TrainingArguments(
+            output_dir="/workspace/checkpoints",
+            per_device_train_batch_size=2,
+            gradient_accumulation_steps=8,
+            num_train_epochs=3,
+            warmup_steps=10,
+            learning_rate=2e-4,
+            weight_decay=0.0,
+            lr_scheduler_type="cosine",
+            optim="adamw_8bit",
+            bf16=is_bf16_supported(),
+            fp16=not is_bf16_supported(),
+            logging_steps=5,
+            save_steps=50,
+            eval_steps=50,
+            save_total_limit=3,
+            report_to="none",
+        ),
+    )
+    log("  trainer initialized - calling train() ...")
+    trainer.train()
+    log("  [OK] training complete")
+
+    # ================================================================
+    # Step 4/4: Push adapter to HF
+    # ================================================================
+    log("\n[4/4] Saving and pushing adapters to HuggingFace ...")
+    os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
+
+    adapter_dir = "/workspace/qwythos-9b-security-adapter"
+    model.save_pretrained(adapter_dir)
+    tokenizer.save_pretrained(adapter_dir)
+    log(f"  adapters saved to {adapter_dir}")
+
+    from huggingface_hub import HfApi, create_repo
+
+    org_repo = "mxguru1/qwythos-9b-security-unsloth"
+    try:
+        create_repo(org_repo, repo_type="model", private=True, exist_ok=True)
+        log(f"  repo ready: {org_repo}")
+    except Exception as e:
+        log(f"  [WARN] create_repo: {e}")
+
+    api = HfApi(token=os.environ.get("HF_TOKEN", ""))
+    try:
+        api.upload_folder(
+            folder_path=adapter_dir,
+            repo_id=org_repo,
+            repo_type="model",
+        )
+        log("  [OK] adapter pushed to HF")
+    except Exception as e:
+        log(f"  [FAIL] upload: {e}")
+        sys.exit(1)
+
+    log("")
+    log("=" * 50)
+    log("COMPLETE")
+    log(f"Adapter: https://huggingface.co/{org_repo}")
+    log("=" * 50)
+
+
+if __name__ == "__main__":
+    main()