modal/tune.py

from __future__ import annotations

import os
import modal

from typing import Any

modal_any: Any = modal

# Modal app groups the remote fine-tuning job.
app = modal_any.App("inner-space-tuner")

# Custom container image with necessary libraries and local dataset code.
image = (
    modal_any.Image.debian_slim()
    .pip_install(
        "torch",
        "transformers>=4.45.0",
        "peft",
        "trl",
        "accelerate",
        "bitsandbytes",
        "datasets",
        "huggingface_hub",
    )
    .add_local_file(
        os.path.join(os.path.dirname(__file__), "dataset.py"),
        "/root/dataset.py",
    )
)

# Volume keeps checkpoints available across Modal runs.
volume = modal_any.Volume.from_name("inner-space-checkpoints", create_if_missing=True)

# Base model stays under the hackathon parameter limit.
MODEL_ID = "openbmb/MiniCPM5-1B-SFT"
ADAPTER_REPO_ID = "build-small-hackathon/inner-space-1b-sft-cbt"


@app.function(
    image=image,
    gpu="A10G",
    timeout=7200,
    volumes={"/checkpoints": volume},
    secrets=[modal_any.Secret.from_name("huggingface-secret")],
)
def train_lora(
    model_card_content: str,
    hf_token: str | None = None,
    repo_id: str | None = None,
):
    """Fine-tunes MiniCPM on six-section CBT reflections using QLoRA."""
    # Remote-only imports are installed inside the Modal container.
    import torch
    from datasets import Dataset
    from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
    from transformers import (
        AutoModelForCausalLM,
        AutoTokenizer,
        BitsAndBytesConfig,
    )
    from trl import SFTConfig, SFTTrainer
    import io
    from huggingface_hub import login, upload_file

    # Import dataset and prompt structure from local modular file.
    from dataset import (
        build_training_prompt,
        get_training_examples,
        get_chat_training_examples,
    )

    # Load the tokenizer for the base model.
    print(f"Loading tokenizer for {MODEL_ID}...")
    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)

    # Synthetic examples capture reports and follow-up coaching turns.
    print("Preparing training dataset...")
    raw_data = get_training_examples()
    chat_data = get_chat_training_examples()

    # Convert report prompt/response pairs into the model chat template.
    formatted_dataset = []
    for item in raw_data:
        messages = [
            {
                "role": "user",
                "content": build_training_prompt(
                    str(item["prompt"]), int(item["distress_level"])
                ),
            },
            {"role": "assistant", "content": item["response"]},
        ]
        text = tokenizer.apply_chat_template(messages, tokenize=False)
        formatted_dataset.append({"text": text})

    # Add multi-turn examples so follow-up coaching does not drift.
    for messages in chat_data:
        text = tokenizer.apply_chat_template(messages, tokenize=False)
        formatted_dataset.append({"text": text})

    # Print the number of training conversations.
    print(f"Prepared {len(formatted_dataset)} total training conversations.")

    # Convert the list of formatted examples into a Hugging Face Dataset.
    dataset = Dataset.from_list(formatted_dataset)

    # QLoRA keeps training feasible on a single A10G.
    print("Configuring QLoRA...")
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16,
        bnb_4bit_use_double_quant=True,
    )

    # Load the base model in 4-bit mode for adapter training.
    print(f"Loading quantized base model {MODEL_ID}...")
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        quantization_config=bnb_config,
        device_map="auto",
    )

    # Set the EOS token ID to the tokenizer's EOS token, and align the PAD token ID to it as well.
    model.config.eos_token_id = tokenizer.eos_token_id
    model.config.pad_token_id = tokenizer.eos_token_id
    if getattr(model, "generation_config", None) is not None:
        model.generation_config.eos_token_id = tokenizer.eos_token_id
        model.generation_config.pad_token_id = tokenizer.eos_token_id

    # Prepare quantized layers for PEFT adapter updates.
    model = prepare_model_for_kbit_training(model)

    # Target attention projections with enough capacity for structured reflections.
    peft_config = LoraConfig(
        r=16,
        lora_alpha=32,
        target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],
        lora_dropout=0.05,
        bias="none",
        task_type="CAUSAL_LM",
    )
    model = get_peft_model(model, peft_config)
    model.print_trainable_parameters()

    # Train long enough to learn report structure and follow-up chat behavior.
    training_args = SFTConfig(
        output_dir="/checkpoints/inner-space-lora",
        per_device_train_batch_size=1,
        gradient_accumulation_steps=4,
        warmup_steps=10,
        max_steps=220,
        learning_rate=2e-4,
        fp16=False,
        bf16=True,
        logging_steps=5,
        save_strategy="steps",
        save_steps=55,
        save_total_limit=2,
        report_to="none",
        dataset_text_field="text",
        max_length=1536,
    )

    # Train only the adapter weights.
    print("Starting fine-tuning job on Modal...")
    trainer = SFTTrainer(
        model=model,
        train_dataset=dataset,
        args=training_args,
    )
    trainer.train()
    print("Fine-tuning completed successfully!")

    # Save the final adapter into the mounted checkpoint volume.
    print("Saving fine-tuned adapter...")
    model.save_pretrained("/checkpoints/inner-space-final")
    tokenizer.save_pretrained("/checkpoints/inner-space-final")

    # Commit the changes to the persistent volume in Modal
    print("Committing checkpoint changes to Modal Volume...")
    volume.commit()

    # Prefer explicit token and repo values when provided.
    if not hf_token:
        hf_token = os.environ.get("HF_TOKEN")
    if not repo_id:
        repo_id = ADAPTER_REPO_ID

    # Publish the adapter when Hub credentials are available.
    if hf_token:
        login(token=hf_token)
        print(
            f"Pushing fine-tuned adapter to Hugging Face Hub repository: {repo_id}..."
        )
        model.push_to_hub(repo_id)
        tokenizer.push_to_hub(repo_id)
        upload_file(
            path_or_fileobj=io.BytesIO(model_card_content.encode("utf-8")),
            path_in_repo="README.md",
            repo_id=repo_id,
            repo_type="model",
            commit_message="Update InnerSpace adapter model card",
        )
        print("Pushed successfully!")
    else:
        print("HF_TOKEN not set or provided. Skipping publishing step.")


@app.local_entrypoint()
def main():
    # Read CARD.md dynamically from the local filesystem during entrypoint execution.
    meta_path = os.path.join(os.path.dirname(__file__), "CARD.md")
    with open(meta_path, "r", encoding="utf-8") as f:
        model_card = f.read()

    # Launch the remote fine-tuning job from local CLI.
    train_lora.remote(model_card_content=model_card)