File size: 2,208 Bytes
17a6842 d5c7535 17a6842 d5c7535 17a6842 d5c7535 95f1e85 17a6842 95f1e85 17a6842 95f1e85 17a6842 95f1e85 17a6842 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 | # /// script
# dependencies = ["trl>=0.12.0", "peft>=0.7.0", "trackio", "datasets", "transformers", "accelerate", "torch"]
# ///
"""SFT on compile-verified C++ curriculum. Designed for Hugging Face Jobs (uv)."""
import os
from datasets import load_dataset
from peft import LoraConfig
from trl import SFTConfig, SFTTrainer
DATASET_ID = os.environ.get("DATASET_ID", "gonzalolinares/cpp-compiler-curriculum")
MODEL_ID = os.environ.get("MODEL_ID", "Qwen/Qwen2.5-1.5B-Instruct")
HUB_MODEL_ID = os.environ.get("HUB_MODEL_ID", "gonzalolinares/qwen25-1.5b-cpp-sft")
OUTPUT_DIR = os.environ.get("OUTPUT_DIR", "qwen25-1.5b-cpp-sft")
def main() -> None:
ds = load_dataset(DATASET_ID, split="train")
# Keep messages format for SFTTrainer chat templates
if "messages" not in ds.column_names:
raise SystemExit(f"Dataset must have 'messages'; got {ds.column_names}")
split = ds.train_test_split(test_size=0.1, seed=42)
trainer = SFTTrainer(
model=MODEL_ID,
train_dataset=split["train"],
eval_dataset=split["test"],
peft_config=LoraConfig(
r=16,
lora_alpha=32,
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
),
args=SFTConfig(
output_dir=OUTPUT_DIR,
num_train_epochs=4,
per_device_train_batch_size=2,
per_device_eval_batch_size=2,
gradient_accumulation_steps=8,
learning_rate=2e-4,
logging_steps=10,
eval_strategy="steps",
eval_steps=40,
save_strategy="epoch",
save_total_limit=1,
max_length=1024,
bf16=True,
push_to_hub=False, # push once at end (avoid Trackio parquet bug mid-save)
hub_model_id=HUB_MODEL_ID,
report_to="none",
),
)
trainer.train()
trainer.model.push_to_hub(HUB_MODEL_ID, private=False)
trainer.processing_class.push_to_hub(HUB_MODEL_ID, private=False)
print(f"Pushed to {HUB_MODEL_ID}")
if __name__ == "__main__":
main()
|