# =========================================================================================
# Fine-tuning script based on https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/Llama3.2_%281B_and_3B%29-Conversational.ipynb
# This script and HF checkpoint are only intended to showcase how to do finetuning in a way compatible with ExecuTorch
# Only 10 steps are done, and quality of the finetuned model is not evaluated
# =========================================================================================

from unsloth import FastLanguageModel
from unsloth.chat_templates import (
    get_chat_template,
    standardize_data_formats,
    standardize_sharegpt,
    train_on_responses_only,
)

from datasets import load_dataset
from trl import SFTConfig, SFTTrainer
from transformers import DataCollatorForSeq2Seq
import torch
import torch.nn as nn

batch_size = 2
learning_rate = 2e-5
gradient_accumulation_steps = 4
max_steps = 10
full_finetuning = True
qat_scheme = "int8-int4"
output_dir = "/tmp/unsloth_example"


model_id = "unsloth/Qwen3-4B"
chat_template = "qwen3"
max_seq_length = 2048
dtype = torch.bfloat16
load_in_4bit = False

################################################################################
# Define model/tokenizer
################################################################################

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name=model_id,
    max_seq_length=max_seq_length,
    dtype=dtype,
    load_in_4bit =load_in_4bit,
    full_finetuning=full_finetuning,
    qat_scheme=qat_scheme,
)
tokenizer = get_chat_template(tokenizer, chat_template = chat_template)

print("MODEL AFTER LOADING")
print(model)
print(model.config)
print(model._torchao_config)

################################################################################
# Process dataset
################################################################################

def formatting_prompts_func(examples):
    convos = examples["conversations"]
    texts = [tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = False) for convo in convos]
    return { "text" : texts, }
dataset = load_dataset("mlabonne/FineTome-100k", split = "train")
dataset = standardize_sharegpt(dataset)
dataset = dataset.map(formatting_prompts_func, batched = True,)

print("DATASET ENTRY")
print(dataset[0])
print("\n\n")

################################################################################
# Define trainer
################################################################################

trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=dataset,
    dataset_text_field="text",
    max_seq_length=max_seq_length,
    packing=False,
    args=SFTConfig(
        per_device_train_batch_size=batch_size,
        gradient_accumulation_steps=gradient_accumulation_steps,
        warmup_steps=5,
        num_train_epochs=1,
        max_steps=max_steps,
        learning_rate=learning_rate,
        logging_steps=1,
        optim="adamw_8bit",
        weight_decay=0.01,
        lr_scheduler_type="linear",
        seed=3407,
        output_dir="outputs",
        report_to="none",
    ),
)

################################################################################
# Do fine tuning
################################################################################
print("DOING FINETUNING")
trainer_stats = trainer.train()
print(f"{trainer_stats.metrics['train_runtime']} seconds used for training.")
print(
    f"{round(trainer_stats.metrics['train_runtime']/60, 2)} minutes used for training."
)

################################################################################
# Save model
################################################################################
model.save_pretrained(output_dir)
tokenizer.save_pretrained(output_dir)

# ################################################################################
# # Convert model to torchao and save to hub
# ################################################################################

from unsloth.models._utils import _convert_torchao_model
_convert_torchao_model(model)

print('MODEL AFTER CONVERT', model)
print("CONFIG", model.config)


################################################################################
# Push converted model to hub
################################################################################
from huggingface_hub import get_token, whoami

def _get_username():
    token = get_token()
    username = whoami(token=token)["name"]
    return username

username = _get_username()
model_name = model_id.split("/")[-1]
save_to = f"{username}/{model_name}-{qat_scheme}-unsloth-v2"
model.push_to_hub(save_to, safe_serialization=False)
tokenizer.push_to_hub(save_to)

################################################################################
# Load converted from hub and inspect
################################################################################
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained(save_to)
print('model', model)
print("model.embed_tokens.weight", model.model.embed_tokens.weight)
print("model.layers[0].self_attn.q_proj.weight", model.model.layers[0].self_attn.q_proj.weight)
print("lm_head.weight", model.lm_head.weight)

git clone https://github.com/meta-pytorch/executorch-examples.git
open executorch-examples/llm/apple/etLLM.xcodeproj