# /// script
# requires-python = ">=3.13"
# dependencies = [
#     "datasets==4.8.5",
#     "huggingface-hub==1.33.0",
#     "torch==2.14.1",
#     "transformers==5.17.0",
#     "trl==1.15.0",
#     "unsloth==2026.11.2",
#     "unsloth-zoo==2026.11.3",
# ]
# ///

import marimo

__generated_with = "0.25.2"
app = marimo.App(width="medium", auto_download=["html"])


@app.cell(hide_code=True)
def _():
    # Let's start!
    import marimo as mo

    mo.md("# Unsloth Finetuning notebook!")
    return (mo,)


@app.cell(hide_code=True)
def _(mo):
    # Definition of mo_exec and mo_exec stream
    from subprocess import (
        run as _run,
        Popen as _Popen,
        PIPE as _PIPE,
        STDOUT as _STDOUT,
    )

    def mo_exec(*process_args, cwd=None):
        cmdline = " ".join(process_args).split(" ")
        result = _run(
            cmdline,
            capture_output=True,
            text=True,
            cwd=cwd,
        )

        result_text = f"""
    **Command**:    _{" ".join(process_args)}_

    **Exit code**:  {result.returncode}
    """

        if result.stdout:
            result_text += f"""

    **stdout**
    ```
    {result.stdout}
    ```
    """
        if result.stderr:
            result_text += f"""

    **stderr**
    ```
    {result.stderr}
    ```
    """
        return mo.md(result_text)

    def mo_stream(*process_args, cwd=None):
        print(f"🚀 Running: {' '.join(process_args)} ...")

        # Run process with real-time output piping
        cmdline = " ".join(process_args).split(" ")
        process = _Popen(
            cmdline,
            stdout=_PIPE,
            stderr=_STDOUT,
            text=True,
            cwd=cwd,
            bufsize=1,
            errors="replace",  # Replace invalid bytes
        )

        # Print lines as they come in so you see live progress
        for line in process.stdout:
            print(line, end="")

        process.wait()

        if process.returncode != 0:
            raise RuntimeError(
                f"❌ Command failed with exit code {process.returncode}"
            )

    return mo_exec, mo_stream


@app.cell
def _(mo_stream):
    # Important setup!
    import os

    # Avoid Unsloth custom modules compilation. Make sure we can run properly
    os.environ["UNSLOTH_FULLGRAPH"] = "0"
    os.environ["UNSLOTH_NO_COMPILE"] = "1"
    os.environ["UNSLOTH_DISABLE_MOE_GROUPED_MM"] = "1"

    # Extra debug messages in case of errors
    os.environ.setdefault("TORCHDYNAMO_VERBOSE", "1")
    os.environ.setdefault("TORCH_LOGS", "+dynamo")

    # Hide HuggingFace download messages
    os.environ.setdefault("HF_XET_HIGH_PERFORMANCE", "1")
    # Set prebuilt llama.cpp location
    llama_cpp_dir = "/marimo/unsloth/llama.cpp"
    os.environ.setdefault("UNSLOTH_LLAMA_CPP_PATH", llama_cpp_dir)

    # Do the ugly. Install Unsloth stuff via CLI
    mo_stream("unsloth install-kernels")

    from huggingface_hub import HfApi

    hf = HfApi()  # HF_TOKEN is in .env
    print("Connected to HuggingFace!")
    return hf, llama_cpp_dir, os


@app.cell
def _(mo_stream):
    # Show required wheels
    mo_stream("unsloth install-kernels --dry-run")
    print(f"""
    These are the packages to be installed via the cell below or the *Packages* section.

    If needed, install these dependencies (found by hand; better the one from the cell above):

    - causal-conv1d: https://github.com/protopia-ai/causal-conv1d/releases/download/v1.6.2.post7/causal_conv1d-1.6.2.post7+cu130torch2.14cxx11abiTRUE-cp313-cp313-linux_x86_64.whl
    - xformers: https://download.pytorch.org/whl/cu130/xformers-0.0.35-py39-none-manylinux_2_28_x86_64.whl
    - flash-attn: https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.10.0/flash_attn-2.8.3%2Bcu130torch2.14-cp313-cp313-linux_x86_64.whl
    """)
    return


@app.cell
def _(mo_stream):
    # Check Unsloth status
    mo_stream("unsloth doctor")
    return


@app.cell
def _(mo_exec):
    mo_exec("env")
    return


@app.cell
def _(mo, os):
    # Show system status
    import sys
    import torch

    # Silly dependency, but ensures we run the cell with the os import which has all the important stuff.
    _unsloth_patch = True if os.getenv("UNSLOTH_PATCHED") == "1" else False

    fa_avail, xf_avail, cc_avail = False, False, False
    try:
        import flash_attn as _fa

        fa_avail = _fa.__version__
    except ImportError:
        pass

    try:
        import xformers as _xf

        xf_avail = _xf.__version__
    except ImportError:
        pass

    try:
        import causal_conv1d as _cc

        cc_avail = _cc.__version__
    except ImportError:
        pass

    mo.md(f"""
    ## System status

    |**Metric**      |**Value**                      |
    |----------------|-------------------------------|
    |Python version  |{sys.version}                  |
    |PyTorch version |{torch.__version__}            |
    |CUDA available  |{torch.cuda.is_available()}    |
    |CUDA version    |{torch.version.cuda}           |
    |GPU             |{torch.cuda.get_device_name(0)}|
    |Flash Attention |{fa_avail}                     |
    |XFormers        |{xf_avail}                     |
    |causal-conv1d   |{cc_avail}                     |
    |Unsloth applied |{_unsloth_patch}               |

    """)
    return (torch,)


@app.cell(hide_code=True)
def _():
    # Define models and datasets
    # Here are all the models and datasets we want to use for our finetuning

    MODELS = [
        "LiquidAI/LFM2.5-24B-A2B",
        "LiquidAI/LFM2.5-8B-A1B",
        "LiquidAI/LFM2.5-2.6B",
        "LiquidAI/LFM3-1.7B",
    ]

    DATASETS = [
        "nvidia/Nemotron-SFT-SWE-v3",
        "empero-ai/gpt-5.6-luna-sft-900x",
    ]
    return DATASETS, MODELS


@app.cell(hide_code=True)
def _(DATASETS, MODELS, mo):
    # Nice dropdown for selection.

    model_selection = mo.ui.dropdown(
        options=MODELS, value=MODELS[0], label="Model to finetune"
    )
    dataset_selection = mo.ui.dropdown(
        options=DATASETS,
        value=DATASETS[0],
        label="Dataset to use for finetuning",
    )

    mo.vstack([model_selection, dataset_selection])
    return dataset_selection, model_selection


@app.cell(hide_code=True)
def _(dataset_selection, mo, model_selection):
    # Fill the required variables for all next sections!
    base_model_name, dataset_name = (
        model_selection.value,
        dataset_selection.value,
    )

    # Definition of the place where the model ends!
    model_base_name = "-".join(base_model_name.split("-")[1:])
    model_name = f"LiquidCode-2.5-{model_base_name}"
    model_repo = f"ItachiSan/{model_name}"
    model_gguf_repo = model_repo + "-GGUF"

    mo.md(f"""
    Selected combination:

    - Base model: {base_model_name}
    - Dataset: {dataset_name}
    - Model name: {model_name}
    - Model repository: {model_repo}
    """)
    return (
        base_model_name,
        dataset_name,
        model_base_name,
        model_gguf_repo,
        model_name,
        model_repo,
    )


@app.cell
def _(mo):
    mo.md("""
    # Setting up crucial parameters!
    """)
    return


@app.function(hide_code=True)
# Function to calculate the training config.
def calculate_training_config(
    vram_gb: float,
    model_params_b: float,
    num_gpus: int = 1,
    target_gbs: int = 128,
    max_seq_length: int = 2048,
    use_4bit_qlora: bool = True,
):
    """
    Calculates safe per-device batch size and gradient accumulation steps.
    """
    print(f"--- Training Configuration Calculator ---")
    print(f"Hardware: {vram_gb} GB VRAM | GPUs: {num_gpus}")
    print(
        f"Model Size: {model_params_b}B parameters | Max Seq Length: {max_seq_length}"
    )
    print(f"Target Global Batch Size (GBS): {target_gbs}\n")

    # 1. Estimate Base Model VRAM Footprint
    if use_4bit_qlora:
        base_model_vram = (
            model_params_b * 0.55
        )  # ~0.55 GB per billion params in 4-bit
    else:
        base_model_vram = (
            model_params_b * 2.0
        )  # ~2 GB per billion params in 16-bit

    overhead_vram = 4.0  # CUDA context, system overhead, LoRA adapters
    available_vram_for_activations = (
        vram_gb - base_model_vram - overhead_vram
    )

    if available_vram_for_activations <= 0:
        return "❌ Error: Model is too large to fit in this VRAM size even with quantization!"

    # 2. Approximate VRAM consumed per sequence at given sequence length
    # (Empirical scaling factor for 4k/8k context transformers)
    vram_per_sequence = (
        (max_seq_length / 4096) * (model_params_b / 24) * 2.5
    )

    # 3. Calculate safe max per-device batch size (leaving a 15% safety buffer)
    safe_vram_pool = available_vram_for_activations * 0.85
    estimated_bs = int(safe_vram_pool / vram_per_sequence)

    # Bound batch size between 1 and 32, rounding to clean numbers if possible
    per_device_bs = max(1, min(estimated_bs, 32))

    # 4. Calculate Gradient Accumulation Steps to match Target GBS
    total_effective_bs_per_step = per_device_bs * num_gpus

    if target_gbs % total_effective_bs_per_step != 0:
        # Round up or find closest clean division
        grad_accum_steps = max(
            1, round(target_gbs / total_effective_bs_per_step)
        )
        actual_gbs = per_device_bs * grad_accum_steps * num_gpus
    else:
        grad_accum_steps = target_gbs // total_effective_bs_per_step
        actual_gbs = target_gbs

    print(f"✅ Recommended Settings:")
    print(f"   per_device_train_batch_size = {per_device_bs}")
    print(f"   gradient_accumulation_steps = {grad_accum_steps}")
    print(
        f"   Resulting Global Batch Size = {actual_gbs} (Target was {target_gbs})"
    )
    print(
        f"   Estimated VRAM utilization: ~{int(base_model_vram + (per_device_bs * vram_per_sequence) + overhead_vram)} GB / {vram_gb} GB\n"
    )

    return per_device_bs, grad_accum_steps


@app.cell(hide_code=True)
def _(mo_exec):
    # GPU check
    mo_exec("nvidia-smi")
    return


@app.cell(hide_code=True)
def _(mo, model_base_name):
    mo.md(f"""
    Planned model: {model_base_name}

    GPU available: see above.

    Use the data for the next calculation.
    """)
    return


@app.cell
def _(model_base_name):
    # --- Example: Your exact setup (96GB VRAM, 24B Model, 4096 Seq Length) ---

    max_seq_length = 2048  # Adjust as needed (e.g., 4096)
    _target_gbs = 128  # Not sure how to get this, got it from Gemini
    _gpu_size_gb = 96.0  # GPU VRAM GB, get it from nvidia-smi block above

    # Detecting model size from the name
    _model_params_b = model_base_name.rstrip("B")
    if "-" in _model_params_b:
        _model_params_b = _model_params_b.split("-")[0].rstrip("B")
    _model_params_b = float(_model_params_b)

    batch_size, gradient_steps = calculate_training_config(
        vram_gb=_gpu_size_gb,
        model_params_b=_model_params_b,
        num_gpus=1,
        target_gbs=_target_gbs,
        max_seq_length=max_seq_length,
        use_4bit_qlora=True,
    )

    # Nice info print
    print(f"""\
    Configured parameters:

    - Total model size: {_model_params_b}
    - Available VRAM: {_gpu_size_gb}
    - Sequence length: {max_seq_length}
    - Batch size: {batch_size}
    - Gradient steps: {gradient_steps}
    """)
    return batch_size, gradient_steps, max_seq_length


@app.cell(hide_code=True)
def _(mo):
    mo.md("""
    # Training the model!
    """)
    return


@app.cell
def _(base_model_name, max_seq_length):
    # Load the model
    from unsloth import FastLanguageModel

    dtype = None  # None for auto-detection
    load_in_4bit = True  # Use 4-bit quantization to save memory

    # 1. Load base model and tokenizer
    base_model, tokenizer = FastLanguageModel.from_pretrained(
        model_name=base_model_name,  # Or your specific base model path
        max_seq_length=max_seq_length,
        dtype=dtype,
        load_in_4bit=load_in_4bit,
    )

    base_model, tokenizer
    return FastLanguageModel, base_model, tokenizer


@app.cell
def _(base_model):
    # Try detecting modules automatically
    detected_modules = set()

    # Simple map of modules to skip. Their bool value is to print the item once.
    _modules_to_skip = {
        # Modules to skip for LFM2 MoE
        "q_layernorm": False,
        "k_layernorm": False,
        "experts": False,
        "gate": False,
    }

    for _layer in base_model.model.layers:
        for _name, _ in _layer.named_modules():
            # The modules are in form
            # sublayer(?) . module
            # If no dot is there, we have nothing to do here.
            if "." not in _name:
                continue

            _, _module_name = _name.split(".", maxsplit=1)
            # Handle cleanly re-detection
            if "." in _module_name:
                _module_name = _module_name.split(".", maxsplit=1)[0]

            if _module_name in _modules_to_skip:
                if _modules_to_skip[_module_name] is False:
                    print("Skip:", _module_name)
                    _modules_to_skip[_module_name] = True
                continue

            if _module_name not in detected_modules:
                print("Found:", _module_name)

            detected_modules.add(_module_name)

    detected_modules
    return (detected_modules,)


@app.cell
def _(FastLanguageModel, base_model, detected_modules, tokenizer):
    # 2. Configure the LoRA model and get the adapters

    # There are hard-coded values found by inspection and chat with AI
    _target_modules = [
        ## Default values for Unsloth, with extra comments from Gemini
        # 1. Standard Attention Projections
        "q_proj",
        "k_proj",
        "v_proj",
        "o_proj",
        # Unused in Lfm2Moe, but keep for safety
        "gate_proj",
        "up_proj",
        "down_proj",
        ## Added extra modules that should properly target MoE
        # 2. Hybrid Short Convolution Layers (Local temporal mixing)
        "in_proj",
        "out_proj",
        "conv",
        # 3. MoE Expert Projections & Router (Layer naming for LFM MoE)
        "ffn_gate_exps",
        "ffn_up_exps",
        "ffn_down_exps",
        "ffn_gate_inp",
        # 4. Dense FFN layers (Present in non-expert blocks)
        # "ffn_gate", "ffn_up", "ffn_down",
        "w1",
        "w2",
        "w3",  # Fixed after checking the values below
    ]
    # Stick also out own modules for safety.
    _target_modules = sorted(
        list(set(_target_modules + list(detected_modules)))
    )

    print("Training following modules:", _target_modules)

    # Now, get the trainable model
    model = FastLanguageModel.get_peft_model(
        base_model,
        target_modules=_target_modules,
        r=32,  # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
        lora_alpha=64,
        lora_dropout=0,  # Optimized to 0 by default
        bias="none",  # Supports "none", "all"
        use_gradient_checkpointing="unsloth",  # Saves VRAM
        random_state=3407,
        use_rslora=False,  # Rank stabilized LoRA support
        loftq_config=None,  # LoftQ
    )

    model, tokenizer
    return (model,)


@app.cell
def _(model):
    _trainable_params = 0
    _total_params = 0

    for _name, _param in model.named_parameters():
        _total_params += _param.numel()
        if _param.requires_grad:
            _trainable_params += _param.numel()

    print(f"Total parameters: {_total_params:,}")
    print(
        f"Trainable parameters: {_trainable_params:,} ({100 * _trainable_params / _total_params:.2f}%)"
    )
    return


@app.cell(hide_code=True)
def _(tokenizer):
    print("Tokenizer template:")
    print(tokenizer.get_chat_template())
    return


@app.cell(hide_code=True)
def _():
    # Shared preprocessing functions
    import json

    def rjl(value):
        """
        Hand-made Recursive JSON Loads.

        This is because AI sucks at writing this.

        JSON allows for
        - iterables: dict, list
        - values: bool, null, int, float, string

        Some datasets have JSON nested in strings and the default json.loads does
        not handle those. This function does.

        Args:
            value: The value to convert. Does nothing unless it is a JSON string.

        Returns:
            The parsed Python object, or the original value if not a JSON string.
        """
        # Handle iterables first. We want to be sure to loop when needed.
        if isinstance(value, list):
            return [rjl(item) for item in value]

        if isinstance(value, dict):
            return {key: rjl(val) for key, val in value.items()}

        # Other simple case: if not string, no need to json.loads
        if not isinstance(value, str):
            return value

        # Here we have a string, let's see if it is a JSON
        try:
            decoded = json.loads(value)
        except (json.JSONDecodeError, TypeError):
            # The string is not a JSON, so return it.
            return value

        # Now we decoded something, but we don't know what.
        # So, let's recurse! This way we handle all cases
        return rjl(decoded)

    return (rjl,)


@app.cell(hide_code=True)
def _(c, rjl):
    # Nemotron specific preprocessing
    def _fix_nemotron_chat(chat):
        fixed_chat = []

        for message in chat:
            fixed_message = message.copy()

            # Ensure tool calls are data and not JSON strings
            fixed_message["tool_calls"] = rjl(message["tool_calls"])

            # If we are an assistant, ensure the tool calls is a list otherwise the template will fail
            if fixed_message["role"] == "assistant":
                if fixed_message["tool_calls"] is None:
                    fixed_message["tool_calls"] = []

            fixed_chat.append(fixed_message)

        return fixed_chat

    def preprocess_nemotron_sft(examples):
        # First, get the proper column.
        chats = examples["messages"]

        # Check if we are batching or not and return the result accordingly
        batched = not isinstance(chats[0], dict)

        return (
            [_fix_nemotron_chat(c) for c in chats]
            if batched
            else _fix_nemotron_chat(c)
        )

    return (preprocess_nemotron_sft,)


@app.cell
def _(tokenizer):
    # Process the dataset if needed
    def generate_prompts_text(
        examples,
        dataset_name: str | None = None,
        dataset_functions: dict[str, any] | None = None,
    ):
        # If we need to do any extra processing, here we do:
        if dataset_name and dataset_functions:
            examples = dataset_functions[dataset_name](examples)

        return {
            "text": tokenizer.apply_chat_template(
                examples,
                # We tokenize later with Unsloth
                tokenize=False,
                # Not sure why. Perhaps check?
                add_generation_prompt=False,
            )
        }

    return (generate_prompts_text,)


@app.cell
def _(model_name):
    dataset_repo = f"ItachiSan/{model_name}-SFT-chats"
    print(f"Dataset repo: {dataset_repo}")
    return (dataset_repo,)


@app.cell
def _(
    dataset_name,
    dataset_repo,
    generate_prompts_text,
    model_name,
    os,
    preprocess_nemotron_sft,
):
    # v2: do evaluation
    from datasets import load_dataset

    # Fetch the dataset online. Disable only the first time we process the dataset.
    _fetch_online = True
    # Whether to push the dataset for refetching. Requires the repo to exist.
    _push_online = True
    # Choose how much to use for evaluation
    eval_split_size = 0.01  # 1%
    # Good number to detect how many processes we can run
    num_proc = int(os.cpu_count() * 0.75) + 1

    if not _fetch_online:
        # Load the dataset, efficiently!
        dataset = load_dataset(dataset_name, split="train", num_proc=num_proc)

        # Format conversations using the model's chat template.
        # Adjust the column name (e.g., "conversations" or "messages")
        # to match your dataset's schema.

        # IMPORTANT CHANGES!
        # We need to drop the other columns, because... we suck
        # The main issue is that, in our patching of the Nemotron dataset, we convert data.
        # Then, pyarrow complains the JSON data we parse to build the text has different shape (duh...)
        # We bypass the problem by dropping all the other columns, because yes.
        chats = dataset.map(
            generate_prompts_text,
            fn_kwargs=dict(
                dataset_name=dataset_name,
                dataset_functions={
                    ## Add special handlers connections in here if needed.
                    # Select proper column and decode JSON strings
                    "nvidia/Nemotron-SFT-SWE-v3": preprocess_nemotron_sft,
                },
            ),
            batched=True,
            remove_columns=dataset.column_names,
            num_proc=num_proc,
        )

        chats_split = chats.train_test_split(
            test_size=0.01,  # 1% for test size can also be an integer for # of rows
            shuffle=True,  # Should always set to True!
            seed=3407,
        )

        if _push_online:
            chats_split.push_to_hub(
                f"ItachiSan/{model_name}-SFT-chats",
                private=True,
                num_proc=num_proc,
            )
    else:
        chats_split = load_dataset(dataset_repo, num_proc=num_proc)

    # Split the dataset
    train_dataset = chats_split["train"]  # Dataset for training
    eval_dataset = chats_split["test"]  # Dataset for evaluation

    train_dataset, eval_dataset
    return eval_dataset, num_proc, train_dataset


@app.cell
def _(train_dataset):
    # Peek into the chats!

    from random import randint as _randint

    _chat_n = _randint(0, train_dataset.num_rows)

    print(f"Example chat (row {_chat_n}):\n{train_dataset[_chat_n]['text']}")
    return


@app.cell
def _(
    batch_size,
    eval_dataset,
    gradient_steps,
    max_seq_length,
    model,
    model_repo,
    num_proc,
    tokenizer,
    torch,
    train_dataset,
):
    # Configure the SFTTrainer
    from trl import SFTTrainer, SFTConfig

    # Using the formula in the tool section, we can calculate the following.
    training_steps = int(
        round((train_dataset.num_rows / (batch_size * gradient_steps)) / 10)
        * 10
    )
    warmup_steps = training_steps * 0.05

    print(
        f"Training {training_steps} steps with batch: {batch_size}, gradient: {gradient_steps}"
    )

    _sft_args = SFTConfig(
        # Gemini recommended params
        # Numbers in comments are earlier ones, now they are calculated via function above.
        # Math is easy: 16x8 is based on GPU performance
        # Then, n_steps are dataset_size / (16x8)
        # Then warmup steps is 5% of the buffer
        per_device_train_batch_size=batch_size,
        gradient_accumulation_steps=gradient_steps,
        max_steps=training_steps,
        warmup_steps=warmup_steps,
        # Log not too frequently (original: 1)
        logging_steps=5,
        # Better on long trainings (original: linear)
        lr_scheduler_type="cosine",
        # Same as before
        learning_rate=2e-4,
        fp16=not torch.cuda.is_bf16_supported(),
        bf16=torch.cuda.is_bf16_supported(),
        optim="adamw_8bit",
        weight_decay=0.01,
        seed=3407,
        output_dir="outputs",
        report_to="none",
        # We do long training. Ensure we can resume.
        save_strategy="steps",
        save_steps=50,
        save_total_limit=5,
        # Store checkpoints on Hugging Face
        push_to_hub=True,
        hub_model_id=model_repo,
        hub_strategy="checkpoint",
        # v2: do eval
        eval_strategy="steps",  # Runs eval every few steps or epochs.
        eval_steps=0.05,  # Evaluate every 5%
        per_device_eval_batch_size=batch_size,  # Ensure we eval quickly
        eval_accumulation_steps=gradient_steps,  # You can increase this include of batch_size
        # TODO: add early stopping: https://unsloth.ai/docs/basics/finetuning-from-last-checkpoint#how-do-i-do-early-stopping
    )

    trainer = SFTTrainer(
        model=model,
        tokenizer=tokenizer,
        dataset_text_field="text",  # Here we point at our single column
        max_seq_length=max_seq_length,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        dataset_num_proc=num_proc,
        packing=False,  # TODO: check wheter packing caused some issues or not
        args=_sft_args,
    )

    trainer
    return (trainer,)


@app.cell(hide_code=True)
def _(mo, model_repo, os):
    # Reload checkpoints if any
    from huggingface_hub import snapshot_download

    repo_id = model_repo
    local_output_dir = "outputs"

    # Automatically pull existing checkpoints from your HF Hub repo back into the sandbox
    try:
        print(
            f"📥 Checking Hugging Face Hub for existing checkpoints in {repo_id}..."
        )
        snapshot_download(
            repo_id=repo_id,
            local_dir=local_output_dir,
            allow_patterns=[
                "checkpoint-*/**",
                "last-checkpoint/**",
                "*.json",
                "*.safetensors",
            ],
        )
        print("✅ Checkpoints synced successfully from Hub.")
    except Exception as e:
        print(f"⚠️ No remote checkpoints found or repository is empty yet: {e}")

    # Detect the checkpoint if any
    existing_checkpoint = False
    if os.path.isdir(local_output_dir):
        existing_checkpoints = [
            os.path.join(local_output_dir, d)
            for d in os.listdir(local_output_dir)
            if d.startswith("checkpoint-") or d == "last-checkpoint"
        ]
        existing_checkpoints.sort(
            key=lambda x: (
                int(x.split("-")[-1])
                if x.startswith("checkpoint")  # covers local checkpoints
                else 42069
            )  # for 'last-checkpoint' we use a very huge number
        )
        if existing_checkpoints:  # pythonic len(list) > 0
            existing_checkpoint = existing_checkpoints[-1]

    mo.md(f"Checkpoint: {existing_checkpoint}")
    return (existing_checkpoint,)


@app.cell
def _(mo):
    mo.md("""
    # DO IT!
    """)
    return


@app.cell
def _(os):
    # Print env vars if needed to check if some things need to change
    for _k, _v in sorted(os.environ.items(), key=lambda k: k[0]):
        print(f"{_k}={_v}")
    return


@app.cell
def _(existing_checkpoint, model, tokenizer, trainer):
    # Finetune the model
    _force_training = False

    _starting_point = False
    if not _force_training:
        _starting_point = existing_checkpoint

    # Do the thing!
    trainer_stats = trainer.train(resume_from_checkpoint=_starting_point)

    lora_output_dir = "/marimo/model/lora"

    # Save the LoRA adapters and tokenizer
    model.save_pretrained(lora_output_dir)
    tokenizer.save_pretrained(lora_output_dir)

    print("Model saved to", lora_output_dir)

    trainer_stats
    return


@app.cell
def _(model, model_repo, tokenizer):
    # Publish all the things!
    model.push_to_hub(model_repo)
    tokenizer.push_to_hub(model_repo)
    model.push_to_hub_merged(
        model_repo,
        tokenizer=tokenizer,
        save_method="merged_16bit",
    )
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md("""
    # Testing the model!
    """)
    return


@app.cell(hide_code=True)
def _(os, torch):
    # If you get TorchDynamo errors like:
    #  File "/tmp/torchinductor_root/r7/cr75u36gyrpc7ppml5xne6rruogr56gjefh5axla5phsc2ix5gjp.py"
    #, line 101, in call
    #    buf2 = torch.ops.unsloth_zoo.routed_nf4_gate_up.default(many_args...)
    #NameError: name 's78' is not defined. Did you mean: 's68'?
    # Run this cell. It clears the corrupted TorchInductor cache and reset dynamo graphs
    import shutil
    shutil.rmtree('/tmp/torchinductor_root', ignore_errors=True)
    torch._dynamo.reset()

    # Also, disable all compilation and Dynamo tracing for inference to avoid shape generation bugs
    os.environ["UNSLOTH_NO_COMPILE"] = "1"
    os.environ["UNSLOTH_FULLGRAPH"] = "0"
    # If truly needed, run also the next line
    #os.environ["TORCH_COMPILE"] = "0"
    return


@app.cell
def _(model_repo, os):
    # If available, load the LoRA directly. Otherwise, we fetch from the repo.
    # Used both for chat and export
    loadable_model_path = model_repo
    _lora_model_path = "/marimo/model/lora"

    if os.path.isfile(_lora_model_path + os.sep + "config.json"):
        loadable_model_path = _lora_model_path

    print("Model to load:", loadable_model_path)
    return (loadable_model_path,)


@app.cell
def _(FastLanguageModel, loadable_model_path):
    # Load base model + adapters directly
    test_model, test_tokenizer = FastLanguageModel.from_pretrained(
        model_name=loadable_model_path,
        max_seq_length=2048,
        dtype=None,
        load_in_4bit=True,
    )

    # Enable fast inference mode
    FastLanguageModel.for_inference(test_model)

    test_model, test_tokenizer
    return test_model, test_tokenizer


@app.cell
def _(mo, test_model, test_tokenizer):
    # Hook our model in a chat box!
    def ui_chat_with_our_model(messages: list[mo.ai.ChatMessage], config: mo.ai.ChatModelConfig):
        # Match the custom messages objects to things we can work with:
        chat_history = [
            {"role": message.role, "content": message.content}
            for message in messages
        ]

        # 1. Get those inputs matched
        inputs = test_tokenizer.apply_chat_template(
            chat_history,
            tokenize=True,
            add_generation_prompt=True,
            return_tensors="pt",
        ).to("cuda")


        # 3. Generate response
        return test_model.generate(
            input_ids=inputs,
            use_cache=True,
            max_new_tokens=config.max_tokens,
            temperature=config.temperature,
            top_p=config.top_p,
            top_k=config.top_k,
            #frequency_penalty=config.frequency_penalty,
            #presence_penalty=config.presence_penalty,
        )

    mo.ui.chat(
        ui_chat_with_our_model,
        show_configuration_controls=True,
    )
    return


@app.cell
def _(mo):
    mo.md("""
    # Export time!
    """)
    return


@app.cell
def _(FastLanguageModel, loadable_model_path):
    # Wire the model for export
    print("🔄 Loading uploaded model weights for export to GGUF...")
    export_model, export_tokenizer = FastLanguageModel.from_pretrained(
        model_name=loadable_model_path,
        max_seq_length=2048,
        dtype=None,
        load_in_4bit=True,
    )
    print(
        "✅ Model loaded! Variables are `export_model` and `export_tokenizer`."
    )
    export_model, export_tokenizer
    return export_model, export_tokenizer


@app.cell(hide_code=True)
def _(mo, os):
    # Define change_dir to allow Unsloth to run properly.

    from contextlib import contextmanager

    @contextmanager
    def change_dir(path):
        prev_dir = os.getcwd()

        if not os.path.isdir(path):
            os.mkdir()

        os.chdir(path)
        try:
            yield
        finally:
            os.chdir(prev_dir)

    mo.md("""
    `change_dir` is now defined.

    Remember to use it for Unsloth GGUF operations as such:
    ```python
    with change_dir("/home/marimo/.unsloth/llama.cpp"):
        do_the_thing.
    ```
    """)
    return


@app.cell(hide_code=True)
def _(manual_model, manual_tokenizer):
    print("🔄 Merging adapters to 16-bit safetensors...")

    # Save the merged model as standard HF format first (bypasses GGUF conversion script)
    manual_model.save_pretrained_merged(
        "merged_model",
        manual_tokenizer,
        save_method="merged_16bit",  # Merges LoRA into base weights
    )

    print("✅ Merged model successfully saved to ./merged_model_16bit!")
    return


@app.cell
def _(llama_cpp_dir):
    # Retrieve Unsloth llama.cpp
    # Original referece here:
    # https://github.com/unslothai/unsloth-zoo/blob/main/unsloth_zoo/llama_cpp.py

    print("🔄 Fetching Unsloth llama.cpp ...")
    from unsloth_zoo.llama_cpp import install_llama_cpp

    install_llama_cpp(
        llama_cpp_folder=llama_cpp_dir, gpu_support=True, print_output=True
    )
    print(f"✅ llama.cpp is now available at {llama_cpp_dir}")
    return


@app.cell
def _(export_model, export_tokenizer):
    # First, export the model as merged bits
    model_merged_dir = "/marimo/model/merged"
    export_model.save_pretrained_merged(
        model_merged_dir, export_tokenizer, save_method="merged_16bit"
    )
    return (model_merged_dir,)


@app.cell
def _(llama_cpp_dir, mo_stream, model_merged_dir, model_name, os):
    # Second, convert to F16 GGUF as first step
    model_gguf_dir = "/marimo/model/gguf"
    os.makedirs(model_gguf_dir, exist_ok=True)

    def model_gguf_file(quant):
        return f"{model_gguf_dir}/{model_name}-{quant}.gguf"

    os.makedirs(model_gguf_dir, exist_ok=True)

    _f16_model_path = model_gguf_file("F16")
    if not os.path.isfile(_f16_model_path):
        mo_stream(
            "python",
            f"{llama_cpp_dir}/convert_hf_to_gguf.py",
            model_merged_dir,
            "--outfile",
            model_gguf_file("F16"),
            "--outtype",
            "f16",
            "--split-max-size",
            "50G",
            cwd=llama_cpp_dir,
        )
    else:
        print("Conversion to F16 GGUF already done!")
    return model_gguf_dir, model_gguf_file


@app.cell
def _(llama_cpp_dir, mo_stream, model_gguf_file, os):
    # Now quantize all the files!
    quantize_quants = [
        "Q2_K",
        "Q3_K_M",
        "Q4_K_M",
        "Q5_K_M",
        "Q6_K",
        "Q8_0",
        "TQ1_0",
        "TQ2_0",
    ]

    for _quant in quantize_quants:
        # Handle special quant file naming
        _model_gguf_file = model_gguf_file(_quant)
        if _quant == "Q2_K":
            _model_gguf_file = _model_gguf_file.replace(_quant, "Q2_K_M")
        elif _quant == "Q6_K":
            _model_gguf_file = _model_gguf_file.replace(_quant, "Q6_K_M")
        elif _quant == "Q8_0":
            _model_gguf_file = _model_gguf_file.replace(_quant, "Q8_K_M")

        # Avoid doing the work again.
        if os.path.isfile(_model_gguf_file):
            print(
                f"✅ File for quant {_quant} already exists: {_model_gguf_file}"
            )
            continue

        print(f"✨ Working on quant: {_quant}")
        mo_stream(
            f"{llama_cpp_dir}/llama-quantize",
            model_gguf_file("F16"),
            _model_gguf_file,
            _quant,  # quant type
            "16",  # number of threads
        )
        print(f"✅ File for quant {_quant} now at: {_model_gguf_file}")
    return


@app.cell
def _(hf, model_gguf_dir, model_gguf_repo):
    # Finally push the stuff!
    print("Starting the upload...")
    hf.upload_folder(
        folder_path=model_gguf_dir,
        repo_id=model_gguf_repo,
        repo_type="model",
        allow_patterns=["*.gguf", "*.json", "*.md", "*.txt"],
    )

    print("All GGUF files uploaded successfully!")
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md("""
    # Experiments zone!
    """)
    return


@app.cell(hide_code=True)
def _(mo):
    mo.md(r"""
    Nothing here for now...
    """)
    return


if __name__ == "__main__":
    app.run()