Upload train_aya_sft.py with huggingface_hub

train_aya_sft.py

@@ -0,0 +1,91 @@
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#     "trl>=0.12.0",
+#     "peft>=0.7.0",
+#     "trackio",
+#     "datasets",
+#     "transformers",
+#     "torch"
+# ]
+# ///
+
+import os
+from datasets import load_dataset
+from peft import LoraConfig
+from trl import SFTTrainer, SFTConfig
+import trackio
+
+def format_chatml(example):
+    # Language map based on subset
+    lang_map = {
+        'Lug_Uga': 'Luganda',
+        'Swa_Ken': 'Kiswahili',
+        'Aka_Gha': 'Akan',
+        'Amh_Eth': 'Amharic',
+        'Eng_Uga': 'English',
+        'Eng_Gha': 'English',
+        'Eng_Eth': 'English',
+        'Eng_Ken': 'English'
+    }
+    lang = lang_map.get(example['subset'], 'the Target Language')
+    
+    system_prompt = f"You are a helpful, medically accurate AI assistant fluent in {lang}. Answer the following health question accurately and completely."
+    
+    messages = [
+        {"role": "system", "content": system_prompt},
+        {"role": "user", "content": str(example['input'])},
+        {"role": "assistant", "content": str(example['output'])}
+    ]
+    return {"messages": messages}
+
+def main():
+    print("Loading dataset...")
+    dataset = load_dataset("Bedru/zindi-multilingual-health-qa")
+    
+    print("Formatting dataset...")
+    train_dataset = dataset['train'].map(format_chatml, remove_columns=dataset['train'].column_names)
+    eval_dataset = dataset['validation'].map(format_chatml, remove_columns=dataset['validation'].column_names)
+    
+    print("Initializing trainer...")
+    trainer = SFTTrainer(
+        model="CohereForAI/aya-23-8B",
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+        peft_config=LoraConfig(
+            r=32, 
+            lora_alpha=64, 
+            target_modules=["q_proj", "v_proj", "k_proj", "o_proj"], 
+            task_type="CAUSAL_LM"
+        ),
+        args=SFTConfig(
+            output_dir="zindi-health-qa-aya",
+            push_to_hub=True,
+            hub_model_id="sujanadh/zindi-health-qa-aya-23",
+            num_train_epochs=2,
+            per_device_train_batch_size=2,
+            gradient_accumulation_steps=8,
+            learning_rate=2e-5,
+            eval_strategy="steps",
+            eval_steps=100,
+            save_strategy="steps",
+            save_steps=100,
+            logging_steps=10,
+            report_to="trackio",
+            project="zindi-health-qa",
+            run_name="aya-23-8b-sft-run-1",
+            bf16=True,
+            max_length=1024,
+            hub_strategy="every_save",
+        )
+    )
+
+    print("Starting training...")
+    trainer.train()
+    
+    print("Pushing final model to hub...")
+    trainer.push_to_hub()
+    print("Training complete!")
+
+if __name__ == "__main__":
+    main()