utrgen/train.py

#!/usr/bin/env python3
# train_lm_adamw_8bit.py
#
# Classic Hugging Face Trainer script for causal LM fine-tuning on:
#   ./train.txt and ./valid.txt
#
# Extra:
# - Replaces literal "\\n" with real newline "\n" inside each sample.
# - Ensures each text sample ends with literal "</s>".
# - Prints “what the model sees” before training.

import os
import json
import random

import torch

from datasets import load_dataset
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    TrainingArguments,
    Trainer,
    set_seed,
)


# ----------------------------
# Paths / knobs
# ----------------------------

TOKENIZER_PATH = "./checkpoint-55000"
POLICY_CKPT = "./checkpoint-55000"

TRAIN_TXT = "./train.txt"
VALID_TXT = "./valid.txt"

OUT_DIR = "./runs/lm_finetune_adamw8bit"
LOG_DIR = os.path.join(OUT_DIR, "tb")

SEED = 42
MAX_LENGTH = 512

EOS_STR = "</s>"

# Preview controls
PREVIEW_TEXT_SAMPLES = 8
PREVIEW_TOKEN_TRUNC = 256
PREVIEW_BATCH_ITEMS = 2
WRITE_PREVIEW_JSONL = True


# Training knobs
MAX_STEPS = 50000
PER_DEVICE_TRAIN_BATCH_SIZE = 16
PER_DEVICE_EVAL_BATCH_SIZE = 64
GRADIENT_ACCUMULATION_STEPS = 1

LEARNING_RATE = 5e-5
WEIGHT_DECAY = 0.0
WARMUP_RATIO = 0.01
LR_SCHEDULER_TYPE = "cosine"

DATALOADER_NUM_WORKERS = 16



def preprocess_text(text: str) -> str:
    """
    1) Convert literal '\\n' into a real newline '\n' inside the sample.
    2) Append literal EOS_STR if missing.
    """
    t = text.rstrip("\n")
    t = t.replace("\\n", "\n")
    if not t.rstrip().endswith(EOS_STR):
        t = t + EOS_STR
    return t


def safe_makedirs(path: str) -> None:
    os.makedirs(path, exist_ok=True)


def tokens_preview(tokenizer, ids, max_tokens=80):
    toks = tokenizer.convert_ids_to_tokens(ids)
    if len(toks) > max_tokens:
        return toks[:max_tokens] + ["…"]
    return toks


def dump_jsonl(path: str, rows):
    with open(path, "w", encoding="utf-8") as f:
        for r in rows:
            f.write(json.dumps(r, ensure_ascii=False) + "\n")


class CausalLMDataCollator:
    """
    Pads input_ids and labels for causal language modeling.
    """

    def __init__(self, tokenizer):
        self.tokenizer = tokenizer

    def __call__(self, features):
        batch = self.tokenizer.pad(
            {"input_ids": [f["input_ids"] for f in features]},
            padding=True,
            return_tensors="pt",
        )

        labels = []
        max_len = batch["input_ids"].shape[1]

        for f in features:
            ids = f["labels"]
            ids = ids + [self.tokenizer.pad_token_id] * (max_len - len(ids))
            labels.append(ids[:max_len])

        batch["labels"] = torch.tensor(labels, dtype=torch.long)
        return batch


def main():

    set_seed(SEED)
    torch.backends.cuda.matmul.allow_tf32 = True

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    print("CUDA available:", torch.cuda.is_available())
    if torch.cuda.is_available():
        print("GPU:", torch.cuda.get_device_name(0))

    # ----------------------------
    # Load tokenizer/model
    # ----------------------------
    tokenizer = AutoTokenizer.from_pretrained(
        TOKENIZER_PATH,
        use_fast=True,
        trust_remote_code=True,
        return_token_type_ids=False,
    )

    if tokenizer.pad_token is None:
        if tokenizer.eos_token is not None:
            tokenizer.pad_token = tokenizer.eos_token
        else:
            tokenizer.add_special_tokens({"pad_token": "<|pad|>"})

    model = AutoModelForCausalLM.from_pretrained(
        POLICY_CKPT,
        torch_dtype=torch.float32,
        device_map=None,
        trust_remote_code=True,
    )

    if len(tokenizer) != model.get_input_embeddings().num_embeddings:
        model.resize_token_embeddings(len(tokenizer))

    model.config.use_cache = False
    model.to(device)

    # ----------------------------
    # Load datasets
    # ----------------------------
    raw = load_dataset(
        "text",
        data_files={"train": TRAIN_TXT, "validation": VALID_TXT},
    )

    raw = raw.filter(lambda x: len(x["text"].strip()) > 0)

    def preprocess_map(batch):
        processed = [preprocess_text(t) for t in batch["text"]]
        return {"text": processed}

    raw = raw.map(
        preprocess_map,
        batched=True,
        desc="Preprocess: literal \\n -> newline + ensure </s>",
    )

    # ----------------------------
        # ----------------------------
    def tokenize_batch(batch):
        texts = batch["text"]

        enc = tokenizer(
            texts,
            truncation=True,
            max_length=MAX_LENGTH,
            padding=False,
        )

        enc["labels"] = [ids[:] for ids in enc["input_ids"]]
        return enc

    tok = raw.map(
        tokenize_batch,
        batched=True,
        remove_columns=[],
        desc="Tokenizing",
    )

    drop_cols = []
    if "text" in tok["train"].column_names:
        drop_cols.append("text")

    train_ds = tok["train"].remove_columns(drop_cols) if drop_cols else tok["train"]
    valid_ds = tok["validation"].remove_columns(drop_cols) if drop_cols else tok["validation"]

    # ----------------------------
    # Preview: sample-level
    # ----------------------------
    safe_makedirs(OUT_DIR)

    train_len = len(tok["train"])
    k = min(PREVIEW_TEXT_SAMPLES, train_len)

    random.seed(SEED)
    preview_indices = random.sample(range(train_len), k=k) if train_len > 0 else []

    preview_rows = []

    print("\n=== PREVIEW: samples after preprocess + tokenization + labels ===")

    for idx in preview_indices:
        ex = tok["train"][idx]

        text = ex["text"]
        input_ids = ex["input_ids"]
        labels = ex["labels"]

        label_ids = labels

        row = {
            "split": "train",
            "index": idx,
            "text": text,
            "len_ids": len(input_ids),
            "input_ids_head": input_ids[: min(len(input_ids), PREVIEW_TOKEN_TRUNC)],
            "tokens_head": tokens_preview(
                tokenizer,
                input_ids[: min(len(input_ids), PREVIEW_TOKEN_TRUNC)],
                PREVIEW_TOKEN_TRUNC,
            ),
            "decoded_full": tokenizer.decode(input_ids, skip_special_tokens=False),
            "decoded_labels": tokenizer.decode(label_ids, skip_special_tokens=False),
        }

        preview_rows.append(row)

        print(f"\n--- sample idx={idx} ---")
        print("TEXT postprocessed:")
        print(text)
        print(f"len(input_ids) = {len(input_ids)}")
        print("TOKENS head:")
        print(row["tokens_head"])
        print("DECODED PROMPT:")
        print(row["decoded_prompt"])
        print("DECODED TARGET:")
        print(row["decoded_labels"])

    if WRITE_PREVIEW_JSONL:
        dump_path = os.path.join(OUT_DIR, "preview_samples.jsonl")
        dump_jsonl(dump_path, preview_rows)
        print(f"\nWrote sample previews to: {dump_path}")

    # ----------------------------
    # Collator + Trainer
    # ----------------------------
    data_collator = CausalLMDataCollator(tokenizer=tokenizer)

    args = TrainingArguments(
        output_dir=OUT_DIR,
        overwrite_output_dir=True,

        max_steps=50000,

        per_device_train_batch_size=PER_DEVICE_TRAIN_BATCH_SIZE,
        per_device_eval_batch_size=PER_DEVICE_EVAL_BATCH_SIZE,
        gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS,

        learning_rate=LEARNING_RATE,
        weight_decay=WEIGHT_DECAY,
        warmup_ratio=WARMUP_RATIO,
        lr_scheduler_type=LR_SCHEDULER_TYPE,

        logging_dir=LOG_DIR,
        logging_steps=1,
        report_to=["tensorboard"],

        eval_strategy="epoch",
        save_strategy="epoch",

        dataloader_num_workers=DATALOADER_NUM_WORKERS,
        remove_unused_columns=False,

        load_best_model_at_end=True,
        metric_for_best_model="eval_loss",
        greater_is_better=False,

        optim="adamw_torch",
    )

    trainer = Trainer(
        model=model,
        args=args,
        train_dataset=train_ds,
        eval_dataset=valid_ds,
        tokenizer=tokenizer,
        data_collator=data_collator,
    )

    # ----------------------------
    # Preview: collated batch
    # ----------------------------
    print("\n=== PREVIEW: a collated training batch after padding + labels ===")

    n_take = min(2, len(train_ds))
    features = [train_ds[i] for i in range(n_take)]
    first_batch = data_collator(features)

    input_ids_b = first_batch["input_ids"].detach().cpu()
    labels_b = first_batch["labels"].detach().cpu()
    attn_b = first_batch["attention_mask"].detach().cpu()

    print("Batch keys:", list(first_batch.keys()))
    print("input_ids shape:", tuple(input_ids_b.shape))
    print("labels shape:", tuple(labels_b.shape))
    print("attention_mask shape:", tuple(attn_b.shape))

    batch_preview_rows = []
    n_show = min(PREVIEW_BATCH_ITEMS, input_ids_b.shape[0])

    for i in range(n_show):
        ids = input_ids_b[i].tolist()
        labels = labels_b[i].tolist()
        label_ids = labels

        row = {
            "batch_item": i,
            "input_len": int(attn_b[i].sum().item()),
            "input_ids_head": ids[: min(len(ids), PREVIEW_TOKEN_TRUNC)],
            "tokens_head": tokens_preview(
                tokenizer,
                ids[: min(len(ids), PREVIEW_TOKEN_TRUNC)],
                PREVIEW_TOKEN_TRUNC,
            ),
            "decoded_input": tokenizer.decode(ids, skip_special_tokens=False),
            "decoded_labels": tokenizer.decode(label_ids, skip_special_tokens=False),
        }

        batch_preview_rows.append(row)

        print(f"\n--- batch item {i} ---")
        print("TOKENS head:")
        print(row["tokens_head"])
        print("DECODED INPUT, specials kept:")
        print(row["decoded_input"])
        print("DECODED LABELS, only learned tokens:")
        print(row["decoded_labels"])

    if WRITE_PREVIEW_JSONL:
        dump_path = os.path.join(OUT_DIR, "preview_first_batch.jsonl")
        dump_jsonl(dump_path, batch_preview_rows)
        print(f"\nWrote batch preview to: {dump_path}")

    # ----------------------------
    # Train / save
    # ----------------------------
    train_result = trainer.train()

    final_dir = os.path.join(OUT_DIR, "final")
    trainer.save_model(final_dir)
    tokenizer.save_pretrained(final_dir)

    metrics = train_result.metrics
    metrics["train_samples"] = len(train_ds)

    trainer.log_metrics("train", metrics)
    trainer.save_metrics("train", metrics)
    trainer.save_state()

    eval_metrics = trainer.evaluate()
    eval_metrics["eval_samples"] = len(valid_ds)

    trainer.log_metrics("eval", eval_metrics)
    trainer.save_metrics("eval", eval_metrics)

    print("\nDone.")
    print(f"Final model saved to: {final_dir}")
    print(f"TensorBoard: tensorboard --logdir {OUT_DIR}")


if __name__ == "__main__":
    main()