# /// 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()