Upload train_verifier.py

train_verifier.py

@@ -1,60 +1,51 @@
-"""
-Train Verifier / Judge (Outcome Reward Model)
-===============================================
-RewardTrainer on Qwen3-4B using preference pairs.
-Dataset: narcolepticchicken/speculative-actions-verifier-pref
-"""
 import torch
 from datasets import load_dataset
 from trl import RewardTrainer, RewardConfig
 from peft import LoraConfig
 
-MODEL = "Qwen/Qwen3-4B"
-DATASET = "narcolepticchicken/speculative-actions-verifier-pref"
-OUTPUT = "narcolepticchicken/speculative-verifier-qwen3-4b"
+HUB_ORG = 'narcolepticchicken'
+MODEL = 'Qwen/Qwen3-4B'
+DATASET = f'{HUB_ORG}/speculative-actions-verifier-pref'
+OUTPUT = f'{HUB_ORG}/speculative-verifier-qwen3-4b'
 
-def main():
-    ds = load_dataset(DATASET)
-    train_ds = ds["train"]
-    eval_ds = ds["test"]
+print('Loading dataset...')
+ds = load_dataset(DATASET)
 
-    peft_config = LoraConfig(
-        r=16,
-        lora_alpha=32,
-        target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],
-        modules_to_save=["score"],
-    )
+print('Configuring LoRA...')
+peft_config = LoraConfig(
+    r=16, lora_alpha=32,
+    target_modules=['q_proj', 'v_proj', 'k_proj', 'o_proj'],
+)
 
-    config = RewardConfig(
-        output_dir="/tmp/verifier-out",
-        hub_model_id=OUTPUT,
-        push_to_hub=True,
-        learning_rate=1e-3,
-        per_device_train_batch_size=2,
-        gradient_accumulation_steps=8,
-        num_train_epochs=2,
-        max_seq_length=4096,
-        bf16=True,
-        gradient_checkpointing=True,
-        logging_strategy="steps",
-        logging_steps=10,
-        logging_first_step=True,
-        disable_tqdm=True,
-        report_to="trackio",
-        run_name="verifier-reward-qwen3-4b",
-    )
+print('Configuring Reward Training...')
+config = RewardConfig(
+    output_dir='/tmp/verifier-out',
+    hub_model_id=OUTPUT,
+    push_to_hub=True,
+    learning_rate=2e-4,
+    per_device_train_batch_size=4,
+    gradient_accumulation_steps=4,
+    num_train_epochs=2,
+    bf16=True,
+    gradient_checkpointing=True,
+    logging_strategy='steps',
+    logging_steps=10,
+    logging_first_step=True,
+    disable_tqdm=True,
+    report_to='trackio',
+    run_name='verifier-reward-qwen3-4b',
+)
 
-    trainer = RewardTrainer(
-        model=MODEL,
-        train_dataset=train_ds,
-        eval_dataset=eval_ds,
-        args=config,
-        peft_config=peft_config,
-    )
+print('Initializing Reward Trainer...')
+trainer = RewardTrainer(
+    model=MODEL,
+    train_dataset=ds['train'],
+    eval_dataset=ds['test'],
+    args=config,
+    peft_config=peft_config,
+)
 
-    trainer.train()
-    trainer.push_to_hub()
-    print("Verifier training complete.")
-
-if __name__ == "__main__":
-    main()
+print('Training verifier...')
+trainer.train()
+trainer.push_to_hub()
+print('Verifier training complete.')