Upload train_proposer.py

train_proposer.py

@@ -1,60 +1,53 @@
-"""
-Train Cheap Proposer (Next-Action Predictor)
-=============================================
-SFT on Qwen3-1.7B to predict next action type given conversation state.
-Dataset: narcolepticchicken/speculative-actions-proposer-sft
-"""
 import torch
 from datasets import load_dataset
 from trl import SFTTrainer, SFTConfig
 from peft import LoraConfig
 
-MODEL = "Qwen/Qwen3-1.7B"
-DATASET = "narcolepticchicken/speculative-actions-proposer-sft"
-OUTPUT = "narcolepticchicken/speculative-proposer-qwen3-1.7b"
+HUB_ORG = 'narcolepticchicken'
+MODEL = 'Qwen/Qwen3-1.7B'
+DATASET = f'{HUB_ORG}/speculative-actions-proposer-sft'
+OUTPUT = f'{HUB_ORG}/speculative-proposer-qwen3-1.7b'
 
-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=["embed_tokens", "lm_head"],
-    )
+print('Configuring LoRA...')
+peft_config = LoraConfig(
+    r=16, lora_alpha=32,
+    target_modules=['q_proj', 'v_proj', 'k_proj', 'o_proj'],
+    modules_to_save=['embed_tokens', 'lm_head'],
+)
 
-    config = SFTConfig(
-        output_dir="/tmp/proposer-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=3,
-        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="proposer-sft-qwen3-1.7b",
-    )
+print('Configuring SFT...')
+config = SFTConfig(
+    output_dir='/tmp/proposer-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='proposer-sft-qwen3-1.7b',
+    dataset_text_field='text',
+)
 
-    trainer = SFTTrainer(
-        model=MODEL,
-        train_dataset=train_ds,
-        eval_dataset=eval_ds,
-        args=config,
-        peft_config=peft_config,
-    )
+print('Initializing trainer...')
+trainer = SFTTrainer(
+    model=MODEL,
+    train_dataset=ds['train'],
+    eval_dataset=ds['test'],
+    args=config,
+    peft_config=peft_config,
+)
 
-    trainer.train()
-    trainer.push_to_hub()
-    print("Proposer training complete.")
-
-if __name__ == "__main__":
-    main()
+print('Training proposer...')
+trainer.train()
+trainer.push_to_hub()
+print('Proposer training complete.')