Upload train_sft_capybara.py with huggingface_hub

train_sft_capybara.py

@@ -0,0 +1,99 @@
+# /// script
+# dependencies = ["trl>=0.28.0", "peft>=0.18.1", "datasets", "trackio"]
+# ///
+
+"""SFT training job that fine-tunes Qwen/Qwen2.5-0.5B on the Capybara instructions."""
+
+from __future__ import annotations
+
+import datetime
+import os
+
+from datasets import load_dataset
+from trl import SFTConfig, SFTTrainer
+
+
+def to_prompt_completion(example: dict) -> dict:
+    """Turn a multi-turn conversation into a single prompt-completion pair."""
+    messages = example.get("messages", [])
+    prompt = None
+    completion = None
+    for message in messages:
+        role = (message.get("role") or "").lower()
+        if role == "user" and prompt is None:
+            prompt = message.get("content", "").strip()
+        elif role == "assistant" and prompt and completion is None:
+            completion = message.get("content", "").strip()
+            break
+    if prompt and completion:
+        return {
+            "text": (
+                "### Instruction:\n"
+                f"{prompt}\n\n"
+                "### Response:\n"
+                f"{completion}"
+            )
+        }
+    return {"text": None}
+
+
+def prepare_dataset() -> tuple:
+    """Load Capybara, keep only valid prompt-completion examples, and split it."""
+    dataset = load_dataset("trl-lib/Capybara", split="train")
+    processed = dataset.map(to_prompt_completion, remove_columns=dataset.column_names)
+    processed = processed.filter(lambda example: example["text"] is not None)
+    split = processed.train_test_split(test_size=0.05, seed=42)
+    return split["train"], split["test"]
+
+
+def main() -> None:
+    hf_token = os.environ.get("HF_TOKEN")
+    print("HF_TOKEN:", hf_token)
+
+    print("Building dataset …")
+    train_dataset, eval_dataset = prepare_dataset()
+    print("Train samples:", len(train_dataset))
+    print("Eval samples:", len(eval_dataset))
+
+    model_name = "Qwen/Qwen2.5-0.5B"
+    run_name = f"capybara-sft-{datetime.datetime.utcnow():%Y%m%d-%H%M%S}"
+
+    config = SFTConfig(
+        output_dir="capybara-sft-output",
+        dataset_text_field="text",
+        report_to="trackio",
+        project="capybara-sft",
+        run_name=run_name,
+        trackio_space_id="trackio",
+        eval_strategy="steps",
+        eval_steps=200,
+        logging_steps=50,
+        logging_dir="capybara-sft-output/logs",
+        save_strategy="steps",
+        save_steps=200,
+        save_total_limit=3,
+        num_train_epochs=2,
+        per_device_train_batch_size=4,
+        per_device_eval_batch_size=4,
+        gradient_accumulation_steps=2,
+        gradient_checkpointing=True,
+        learning_rate=2e-5,
+        push_to_hub=True,
+        hub_model_id="davidsmts/sft-capybara-demo",
+        hub_strategy="every_save",
+    )
+
+    trainer = SFTTrainer(
+        model=model_name,
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+        args=config,
+    )
+
+    trainer.train()
+    trainer.push_to_hub()
+
+
+if __name__ == "__main__":
+    main()
+