train_cipher_nano.py

"""
Full fine-tune (not LoRA) of h2oai/h2o-danube3-500m-chat for email triage.

VARIANT EXPERIMENT -- candidate replacement for cipher-nano. SmolLM2 (49K
vocab, 135M/360M) shrinks to the right disk size but plateaus at weak
category/importance accuracy after multiple tuning attempts (LoRA vs full-FT,
epoch sweeps, data reshaping) -- a base-pretraining-quality ceiling, not a
tuning problem (see DEPLOYMENT.md). Qwen2.5-0.5B has the opposite problem:
strong base quality but a 151,936-token vocabulary that floors its disk size
around 340-400MB regardless of quantization, so it can't shrink into nano's
target range either.

Danube3-500M is a plain LlamaForCausalLM with a 32,000-token vocabulary --
much smaller than Qwen/Gemma, comparable to SmolLM2 -- while coming from a
more conventional larger-scale pretraining recipe (h2oai's Danube series).
Worth testing whether it breaks the small-vocab-means-weak-base pattern.

Same data/format/eval as the other nano candidates -- only the base model
differs.

Outputs:
    outputs/danube3-500m-full/model/    - full fine-tuned HF model

Usage:
    python train/train_danube3_500m_full.py
    python train/train_danube3_500m_full.py --epochs 3 --output_dir ./my_run
"""

import argparse
import inspect
import re
from pathlib import Path


def parse_args():
    parser = argparse.ArgumentParser(description="Full fine-tune Danube3-500M for email triage")
    parser.add_argument("--model_name", default="h2oai/h2o-danube3-500m-chat", help="Base HF model")
    parser.add_argument("--train_file", default="train.jsonl", help="Training JSONL")
    parser.add_argument("--val_file", default="val.jsonl", help="Validation JSONL")
    parser.add_argument("--output_dir", default="outputs/danube3-500m-full", help="Root output directory")
    parser.add_argument("--max_seq_length", type=int, default=2048)
    parser.add_argument("--epochs", type=int, default=3)
    parser.add_argument("--lr", type=float, default=5e-5)
    parser.add_argument("--per_device_batch", type=int, default=2)
    parser.add_argument("--gradient_accumulation", type=int, default=4)
    parser.add_argument("--warmup_ratio", type=float, default=0.1)
    parser.add_argument("--seed", type=int, default=3407)
    parser.add_argument("--packing", action="store_true", default=False, help="Pack multiple short examples per sequence (default on)")
    parser.add_argument("--no-packing", dest="packing", action="store_false")
    return parser.parse_args()


def main(args):
    from datasets import disable_caching, load_dataset
    from trl import SFTConfig, SFTTrainer
    from unsloth import FastLanguageModel, is_bfloat16_supported

    disable_caching()

    out_root = Path(args.output_dir)
    model_dir = out_root / "model"
    out_root.mkdir(parents=True, exist_ok=True)

    print(f"Loading {args.model_name} for FULL fine-tune (no LoRA, no quantization) ...")
    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name=args.model_name,
        max_seq_length=args.max_seq_length,
        dtype=None,
        load_in_4bit=False,
        full_finetuning=True,
    )

    print(f"Loading datasets: {args.train_file}, {args.val_file}")
    train_ds = load_dataset("json", data_files=args.train_file, split="train")
    val_ds = load_dataset("json", data_files=args.val_file, split="train")

    def format_chat(example):
        # Danube3-500m-chat uses its own native format -- <|prompt|>...eos
        # for user turns, <|answer|>...eos for assistant turns, strictly
        # alternating, no system role (confirmed against the tokenizer's own
        # chat_template and vocab: ChatML tokens aren't even present).
        # Fold the system prompt into the first user turn's content.
        msgs = example["messages"]
        system_content = ""
        if msgs and msgs[0]["role"] == "system":
            system_content = msgs[0]["content"] + "\n\n"
            msgs = msgs[1:]

        parts = []
        first_user = True
        for msg in msgs:
            content = msg["content"]
            if msg["role"] == "user" and first_user:
                content = system_content + content
                first_user = False
            if msg["role"] == "user":
                parts.append(f"<|prompt|>{content.strip()}{tokenizer.eos_token}")
            else:
                parts.append(f"<|answer|>{content.strip()}{tokenizer.eos_token}")
        text = "".join(parts)
        return {"text": text}

    train_ds = train_ds.map(format_chat, remove_columns=train_ds.column_names)
    val_ds = val_ds.map(format_chat, remove_columns=val_ds.column_names)

    print(f"Train examples: {len(train_ds)}  Validation examples: {len(val_ds)}")

    config_params = inspect.signature(SFTConfig).parameters
    training_kwargs = dict(
        output_dir=str(model_dir),
        num_train_epochs=args.epochs,
        per_device_train_batch_size=args.per_device_batch,
        per_device_eval_batch_size=args.per_device_batch,
        gradient_accumulation_steps=args.gradient_accumulation,
        learning_rate=args.lr,
        warmup_ratio=args.warmup_ratio,
        lr_scheduler_type="cosine",
        optim="adamw_8bit",
        eval_steps=100,
        save_strategy="steps",
        save_steps=100,
        logging_steps=10,
        seed=args.seed,
        fp16=not is_bfloat16_supported(),
        bf16=is_bfloat16_supported(),
        load_best_model_at_end=True,
        metric_for_best_model="eval_loss",
        greater_is_better=False,
        report_to="none",
        dataset_text_field="text",
        packing=args.packing,
    )

    if "eval_strategy" in config_params:
        training_kwargs["eval_strategy"] = "steps"
    else:
        training_kwargs["evaluation_strategy"] = "steps"
    if "max_length" in config_params:
        training_kwargs["max_length"] = args.max_seq_length
    else:
        training_kwargs["max_seq_length"] = args.max_seq_length
    training_args = SFTConfig(**training_kwargs)

    trainer_kwargs = dict(
        model=model,
        train_dataset=train_ds,
        eval_dataset=val_ds,
        args=training_args,
    )
    trainer_params = inspect.signature(SFTTrainer).parameters
    if "processing_class" in trainer_params:
        trainer_kwargs["processing_class"] = tokenizer
    else:
        trainer_kwargs["tokenizer"] = tokenizer

    _orig_convert_tokens_to_ids = tokenizer.convert_tokens_to_ids
    _sentinel_re = re.compile(r"^<([A-Z]+)_TOKEN>$")

    def _convert_tokens_to_ids_patched(token):
        match = _sentinel_re.match(token) if isinstance(token, str) else None
        if match:
            real_id = getattr(tokenizer, f"{match.group(1).lower()}_token_id", None)
            if real_id is not None:
                return real_id
        return _orig_convert_tokens_to_ids(token)

    _orig_prepare_dataset = SFTTrainer._prepare_dataset

    def _prepare_dataset_patched(self, dataset, processing_class, ds_args, *rest, **kw):
        ds_args.dataset_num_proc = None
        return _orig_prepare_dataset(self, dataset, processing_class, ds_args, *rest, **kw)

    SFTTrainer._prepare_dataset = _prepare_dataset_patched

    tokenizer.convert_tokens_to_ids = _convert_tokens_to_ids_patched
    try:
        trainer = SFTTrainer(**trainer_kwargs)
    finally:
        tokenizer.convert_tokens_to_ids = _orig_convert_tokens_to_ids
        SFTTrainer._prepare_dataset = _orig_prepare_dataset

    print("Starting training...")
    trainer.train()

    print(f"Saving full fine-tuned model to {model_dir}")
    model.save_pretrained(model_dir)
    tokenizer.save_pretrained(model_dir)

    print("Done.")


if __name__ == "__main__":
    args = parse_args()
    main(args)