Upload train_unesco_tagger.py with huggingface_hub

train_unesco_tagger.py

@@ -1,7 +1,6 @@
 # /// script
 # dependencies = [
 #     "trl>=0.12.0",
-#     "peft>=0.7.0",
 #     "transformers>=4.36.0",
 #     "accelerate>=0.24.0",
 #     "trackio",
@@ -9,7 +8,6 @@
 # ///
 
 from datasets import load_dataset
-from peft import LoraConfig
 from trl import SFTTrainer, SFTConfig
 
 print("Loading dataset...")
@@ -20,15 +18,15 @@ eval_dataset = dataset["validation"]
 print(f"Train: {len(train_dataset)}, Eval: {len(eval_dataset)}")
 
 config = SFTConfig(
-    output_dir="qwen2.5-3b-unesco-tagger",
+    output_dir="lfm2.5-1.2b-unesco-tagger",
     push_to_hub=True,
-    hub_model_id="unesco-data-ai/qwen2.5-3b-unesco-tagger-v1",
+    hub_model_id="unesco-data-ai/lfm2.5-1.2b-unesco-tagger-v1",
     hub_strategy="every_save",
     num_train_epochs=3,
-    per_device_train_batch_size=2,
-    gradient_accumulation_steps=8,
+    per_device_train_batch_size=4,
+    gradient_accumulation_steps=4,
     learning_rate=2e-5,
-    max_length=2048,
+    max_length=1024,
     logging_steps=10,
     save_strategy="steps",
     save_steps=200,
@@ -37,27 +35,18 @@ config = SFTConfig(
     eval_steps=200,
     warmup_ratio=0.1,
     lr_scheduler_type="cosine",
+    bf16=True,
     report_to="trackio",
     project="unesco-keyword-extraction",
-    run_name="qwen2.5-3b-sft-v1",
-)
-
-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"],
+    run_name="lfm2.5-1.2b-sft-v1",
 )
 
 print("Initializing trainer...")
 trainer = SFTTrainer(
-    model="Qwen/Qwen2.5-3B-Instruct",
+    model="LiquidAI/LFM2.5-1.2B-Instruct",
     train_dataset=train_dataset,
     eval_dataset=eval_dataset,
     args=config,
-    peft_config=peft_config,
 )
 
 print("Starting training...")
@@ -66,4 +55,4 @@ trainer.train()
 print("Pushing to Hub...")
 trainer.push_to_hub()
 
-print("Complete!")
+print("Complete! Model at: https://huggingface.co/unesco-data-ai/lfm2.5-1.2b-unesco-tagger-v1")