#!/root/venvs/train/bin/python
"""
Fine-tuning for QwQ model
"""

# Unsloth should be imported before transformers to ensure all optimizations are applied.
from unsloth import FastLanguageModel, is_bfloat16_supported

from dataclasses import dataclass
from pathlib import Path
from transformers import AutoTokenizer, TrainingArguments
from trl import SFTTrainer
from typing import Optional, List
import argparse
import json
import os

from .dataset import Dataset

@dataclass
class Args:
    def __init__(self, args: Optional[List[str]] = None):
        parser = argparse.ArgumentParser(description='Train SIA model using QwQ')
        parser.add_argument(
            '--config',
            type=Path,
            default=Path('/root/sia/training/config.yaml'),
            help='Path to config file'
        )
        parser.add_argument(
            '--base-model',
            type=str,
            default='unsloth/QwQ-32B-bnb-4bit',
            help='HuggingFace model ID for base model'
        )
        parser.add_argument(
            '--output-dir',
            type=Path,
            required=True,
            help='Directory to save the trained model'
        )
        parser.add_argument(
            '--api-key',
            type=str,
            default=os.environ.get('SIA_HF_API_KEY'),
            help='HuggingFace API key'
        )
        if args is None:
            self.args = parser.parse_args()
        else:
            self.args = parser.parse_args(args)

    @property
    def config_path(self) -> Path:
        return self.args.config

    @property
    def base_model(self) -> str:
        return self.args.base_model
    
    @property
    def output_dir(self) -> Path:
        return self.args.output_dir
    
    @property
    def api_key(self) -> str:
        return self.args.api_key

def main():
    args = Args()
    dataset = Dataset(args.config_path)
    dataset.validate()

    with open('/root/sia/qwq_tokenizer_config.json', 'r') as f:
        tokenizer_config = json.load(f)

    tokenizer = AutoTokenizer.from_pretrained(
        args.base_model,
        **tokenizer_config,
    )

    model, _returned_tokenizer = FastLanguageModel.from_pretrained(
        model_name = args.base_model,
        gpu_memory_utilization = 0.5,
    )

    model = FastLanguageModel.get_peft_model(
        model,
        target_modules = [
            "q_proj", "k_proj", "v_proj", "o_proj",
            "gate_proj", "up_proj", "down_proj",
        ], # Remove QKVO if out of memory
        lora_dropout = 0, # Supports any, but = 0 is optimized
        bias = "none",    # Supports any, but = "none" is optimized
        use_gradient_checkpointing = "unsloth", # True or "unsloth" for very long context
        random_state = 3407,
    )

    training_args = TrainingArguments(
        output_dir=str(args.output_dir) + "_train",
        num_train_epochs=3,
        per_device_train_batch_size=1,
        gradient_accumulation_steps=16,
        gradient_checkpointing=True,
        learning_rate=2e-5,
        lr_scheduler_type="cosine",
        warmup_ratio=0.05,
        weight_decay=0.01,
        fp16=not is_bfloat16_supported(),
        bf16=is_bfloat16_supported(),
        logging_steps=10,
        save_steps=200,
        save_total_limit=3,
        report_to="none",
        optim="adamw_8bit",
    )

    trainer = SFTTrainer(
        model=model,
        tokenizer=tokenizer,
        args=training_args,
        train_dataset=dataset.to_transformers_dataset(tokenizer),
        dataset_text_field="messages",
    )

    trainer.train()
    
    model.save_pretrained_merged(
        str(args.output_dir), 
        tokenizer=tokenizer,
        save_method="merged_4bit_forced"
    )

if __name__ == "__main__":
    main()