ml/scripts/train_dapt.py

"""
Domain-Adaptive Pre-Training (DAPT) for DistilBERT.

Continues masked language model (MLM) pre-training on unlabeled mental health
Reddit text, adapting the encoder to the domain before fine-tuning.

Reference: Gururangan et al. (ACL 2020) — "Don't Stop Pretraining"

Usage:
    python train_dapt.py [options]

Options:
    --model-name: Base model (default: distilbert-base-uncased)
    --corpus: Path to DAPT corpus text file
    --epochs: Number of MLM epochs (default: 3)
    --batch-size: Batch size (default: 32)
    --lr: Learning rate (default: 5e-5)
    --max-length: Max token length (default: 128)
    --output-dir: Where to save DAPT'd model
"""

import argparse
import logging
import math
import os
from pathlib import Path

import torch
from torch.utils.data import DataLoader, Dataset
from tqdm import tqdm
from transformers import (
    AutoModelForMaskedLM,
    AutoTokenizer,
    DataCollatorForLanguageModeling,
    get_linear_schedule_with_warmup,
)

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

os.environ["TOKENIZERS_PARALLELISM"] = "false"


class TextDataset(Dataset):
    """Simple dataset that tokenizes text lines for MLM."""

    def __init__(self, texts: list[str], tokenizer, max_length: int = 128):
        self.encodings = []
        for text in tqdm(texts, desc="Tokenizing"):
            enc = tokenizer(
                text,
                truncation=True,
                max_length=max_length,
                return_tensors="pt",
            )
            self.encodings.append(
                {
                    "input_ids": enc["input_ids"].squeeze(),
                    "attention_mask": enc["attention_mask"].squeeze(),
                }
            )

    def __len__(self):
        return len(self.encodings)

    def __getitem__(self, idx):
        return self.encodings[idx]


def load_corpus(corpus_path: Path, min_length: int = 20) -> list[str]:
    """Load corpus from text file. Split on double-newlines (one post per chunk)."""
    with open(corpus_path) as f:
        raw = f.read()

    # Split on double newlines (each post is separated by \n\n)
    chunks = [c.strip() for c in raw.split("\n\n") if len(c.strip()) >= min_length]
    logger.info(f"Loaded {len(chunks)} text chunks from {corpus_path}")
    return chunks


def evaluate_perplexity(model, dataloader, device) -> float:
    """Compute perplexity on a dataset (lower = better domain fit)."""
    model.eval()
    total_loss = 0
    total_batches = 0

    with torch.no_grad():
        for batch in dataloader:
            input_ids = batch["input_ids"].to(device)
            attention_mask = batch["attention_mask"].to(device)
            labels = batch["labels"].to(device)

            outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
            total_loss += outputs.loss.item()
            total_batches += 1

    avg_loss = total_loss / max(total_batches, 1)
    perplexity = math.exp(avg_loss)
    return perplexity


def main():
    parser = argparse.ArgumentParser(description="Domain-Adaptive Pre-Training (DAPT)")
    parser.add_argument("--model-name", type=str, default="distilbert-base-uncased")
    parser.add_argument("--corpus", type=str, default=None, help="Path to corpus .txt file")
    parser.add_argument("--epochs", type=int, default=3)
    parser.add_argument("--batch-size", type=int, default=32)
    parser.add_argument("--lr", type=float, default=5e-5)
    parser.add_argument("--max-length", type=int, default=128)
    parser.add_argument("--mlm-probability", type=float, default=0.15)
    parser.add_argument("--output-dir", type=str, default=None)
    parser.add_argument("--eval-split", type=float, default=0.05, help="Fraction for eval (default: 5%)")
    args = parser.parse_args()

    base_dir = Path(__file__).parent.parent
    corpus_path = Path(args.corpus) if args.corpus else base_dir / "data" / "dapt_corpus" / "dapt_corpus.txt"
    output_dir = Path(args.output_dir) if args.output_dir else base_dir / "models" / "v2_dapt_base"
    output_dir.mkdir(parents=True, exist_ok=True)

    # Device
    if torch.backends.mps.is_available():
        device = torch.device("mps")
    elif torch.cuda.is_available():
        device = torch.device("cuda")
    else:
        device = torch.device("cpu")
    logger.info(f"Using device: {device}")

    # Load corpus
    logger.info(f"Loading corpus from {corpus_path}...")
    texts = load_corpus(corpus_path)
    if not texts:
        logger.error("No texts loaded. Run collect_dapt_corpus.py first.")
        return

    # Split into train/eval
    n_eval = max(int(len(texts) * args.eval_split), 100)
    eval_texts = texts[:n_eval]
    train_texts = texts[n_eval:]
    logger.info(f"Train: {len(train_texts)} chunks, Eval: {len(eval_texts)} chunks")

    # Tokenizer + Model (MLM head)
    logger.info(f"Loading {args.model_name} for masked language modeling...")
    tokenizer = AutoTokenizer.from_pretrained(args.model_name)
    model = AutoModelForMaskedLM.from_pretrained(args.model_name)
    model.to(device)

    # Datasets
    logger.info("Building datasets...")
    train_dataset = TextDataset(train_texts, tokenizer, args.max_length)
    eval_dataset = TextDataset(eval_texts, tokenizer, args.max_length)

    # MLM data collator — handles masking 15% of tokens
    data_collator = DataCollatorForLanguageModeling(
        tokenizer=tokenizer,
        mlm=True,
        mlm_probability=args.mlm_probability,
    )

    num_workers = 0 if device.type == "mps" else 2
    train_loader = DataLoader(
        train_dataset,
        batch_size=args.batch_size,
        shuffle=True,
        collate_fn=data_collator,
        num_workers=num_workers,
    )
    eval_loader = DataLoader(
        eval_dataset,
        batch_size=args.batch_size,
        collate_fn=data_collator,
        num_workers=num_workers,
    )

    # Optimizer + scheduler
    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr)
    total_steps = len(train_loader) * args.epochs
    scheduler = get_linear_schedule_with_warmup(
        optimizer,
        num_warmup_steps=total_steps // 10,
        num_training_steps=total_steps,
    )

    # Measure perplexity BEFORE DAPT
    logger.info("Measuring pre-DAPT perplexity on domain text...")
    pre_dapt_ppl = evaluate_perplexity(model, eval_loader, device)
    logger.info(f"Pre-DAPT perplexity: {pre_dapt_ppl:.2f}")

    # Training loop
    logger.info(f"\nStarting DAPT: {args.epochs} epochs, {len(train_loader)} batches/epoch")
    best_ppl = pre_dapt_ppl

    for epoch in range(args.epochs):
        model.train()
        total_loss = 0
        n_batches = 0

        for batch in tqdm(train_loader, desc=f"DAPT Epoch {epoch + 1}/{args.epochs}"):
            input_ids = batch["input_ids"].to(device)
            attention_mask = batch["attention_mask"].to(device)
            labels = batch["labels"].to(device)

            outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
            loss = outputs.loss

            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            scheduler.step()
            optimizer.zero_grad()

            total_loss += loss.item()
            n_batches += 1

        avg_train_loss = total_loss / n_batches
        eval_ppl = evaluate_perplexity(model, eval_loader, device)

        logger.info(
            f"Epoch {epoch + 1}: train_loss={avg_train_loss:.4f}, "
            f"eval_perplexity={eval_ppl:.2f} "
            f"(pre-DAPT was {pre_dapt_ppl:.2f})"
        )

        if eval_ppl < best_ppl:
            best_ppl = eval_ppl
            # Save best model
            model.save_pretrained(output_dir)
            tokenizer.save_pretrained(output_dir)
            logger.info(f"  Saved best DAPT model (ppl={best_ppl:.2f})")

    # Final save if never saved (all epochs worse than pre-DAPT, unlikely)
    if not (output_dir / "config.json").exists():
        model.save_pretrained(output_dir)
        tokenizer.save_pretrained(output_dir)

    # Save DAPT metadata
    import json

    dapt_meta = {
        "base_model": args.model_name,
        "epochs": args.epochs,
        "batch_size": args.batch_size,
        "learning_rate": args.lr,
        "max_length": args.max_length,
        "mlm_probability": args.mlm_probability,
        "train_chunks": len(train_texts),
        "eval_chunks": len(eval_texts),
        "pre_dapt_perplexity": pre_dapt_ppl,
        "post_dapt_perplexity": best_ppl,
        "perplexity_reduction": f"{((pre_dapt_ppl - best_ppl) / pre_dapt_ppl) * 100:.1f}%",
    }
    with open(output_dir / "dapt_metadata.json", "w") as f:
        json.dump(dapt_meta, f, indent=2)

    print(f"\n{'=' * 60}")
    print("DAPT COMPLETE")
    print(f"{'=' * 60}")
    print(f"Pre-DAPT perplexity:  {pre_dapt_ppl:.2f}")
    print(f"Post-DAPT perplexity: {best_ppl:.2f}")
    print(f"Reduction:            {((pre_dapt_ppl - best_ppl) / pre_dapt_ppl) * 100:.1f}%")
    print(f"\nDAPT model saved to: {output_dir}")
    print("Use this as base model for fine-tuning:")
    print(f"  python train_redsm5_model.py --model-name {output_dir}")


if __name__ == "__main__":
    main()