src/train.py

"""
Fine-tune google/muril-base-cased for binary scam classification with HuggingFace Trainer.

TrainingArguments uses ``evaluation_strategy`` (Transformers 4.40); targets: F1 > 0.88,
precision > 0.90, recall > 0.85. If ``eval_loss`` rises while F1 falls, reduce epochs or increase early stopping.
"""

from __future__ import annotations

import json
import logging
import sys
from dataclasses import dataclass
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F
import evaluate
from datasets import DatasetDict, load_from_disk
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    TrainingArguments,
    Trainer,
    EarlyStoppingCallback,
)

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
log = logging.getLogger(__name__)

ROOT = Path(__file__).resolve().parent.parent
TOKEN_DIR = ROOT / "data" / "processed" / "tokenized"
SAVED_TOKENIZER_DIR = TOKEN_DIR / "muril_tokenizer"
MODEL_OUT = ROOT / "models" / "muril_scam_v1"
LOGS_DIR = ROOT / "models" / "logs"

MODEL_NAME = "google/muril-base-cased"
NUM_LABELS = 2
LABEL2ID = {"safe": 0, "scam": 1}
ID2LABEL = {0: "safe", 1: "scam"}

CUSTOM_TOKENS = ["[URL]", "[PHONE]", "[EMAIL]", "[AMOUNT]", "[CODE]", "[AADHAAR]", "[PAN]"]


@dataclass
class TrainConfig:
    num_epochs: int = 5
    train_batch_size: int = 16
    eval_batch_size: int = 32
    learning_rate: float = 2e-5
    warmup_ratio: float = 0.1
    weight_decay: float = 0.01
    max_grad_norm: float = 1.0
    seed: int = 42
    fp16: bool = False
    early_stopping_patience: int = 2
    logging_steps: int = 50
    save_total_limit: int = 2

    def __post_init__(self) -> None:
        object.__setattr__(self, "fp16", torch.cuda.is_available())


CFG = TrainConfig()


class WeightedTrainer(Trainer):
    def __init__(self, class_weights: list[float], *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._ce_weight = torch.tensor(class_weights, dtype=torch.float32)
        w0, w1 = class_weights[0], class_weights[1]
        log.info("WeightedTrainer | CrossEntropy weights: safe=%.3f scam=%.3f", w0, w1)

    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
        labels = inputs.pop("labels")
        outputs = model(**inputs)
        logits = outputs.logits
        weight = self._ce_weight.to(device=logits.device, dtype=logits.dtype)
        loss = F.cross_entropy(logits, labels, weight=weight)
        return (loss, outputs) if return_outputs else loss


def build_compute_metrics():
    accuracy_metric = evaluate.load("accuracy")
    f1_metric = evaluate.load("f1")
    precision_metric = evaluate.load("precision")
    recall_metric = evaluate.load("recall")

    def compute_metrics(eval_pred):
        logits, labels = eval_pred
        logits_t = torch.as_tensor(logits)
        probs = torch.softmax(logits_t, dim=-1).numpy()
        preds = np.argmax(logits, axis=-1)
        scam_prob = probs[:, 1]

        acc = accuracy_metric.compute(predictions=preds, references=labels)
        f1_bin = f1_metric.compute(
            predictions=preds, references=labels, average="binary"
        )
        prec = precision_metric.compute(
            predictions=preds, references=labels, average="binary"
        )
        rec = recall_metric.compute(
            predictions=preds, references=labels, average="binary"
        )

        try:
            from sklearn.metrics import roc_auc_score

            auc = float(roc_auc_score(labels, scam_prob))
        except Exception:
            auc = 0.0

        return {
            "accuracy": round(acc["accuracy"], 4),
            "f1": round(f1_bin["f1"], 4),
            "precision": round(prec["precision"], 4),
            "recall": round(rec["recall"], 4),
            "auc_roc": round(auc, 4),
        }

    return compute_metrics


def load_artifacts() -> tuple[DatasetDict, list[float], AutoTokenizer]:
    if not TOKEN_DIR.exists():
        raise FileNotFoundError(
            f"Tokenized data missing: {TOKEN_DIR}\nRun: python src/preprocess.py",
        )

    log.info("Loading dataset from %s", TOKEN_DIR)
    dataset = load_from_disk(str(TOKEN_DIR))
    log.info("  Splits: %s", {k: len(v) for k, v in dataset.items()})

    weights_path = TOKEN_DIR / "class_weights.json"
    if weights_path.exists():
        with open(weights_path, encoding="utf-8") as f:
            class_weights = json.load(f)["weights"]
        log.info("  Class weights: %s", class_weights)
    else:
        log.warning("class_weights.json missing; using [1.0, 1.0]")
        class_weights = [1.0, 1.0]

    if SAVED_TOKENIZER_DIR.exists():
        log.info("Loading tokenizer from %s", SAVED_TOKENIZER_DIR)
        tokenizer = AutoTokenizer.from_pretrained(str(SAVED_TOKENIZER_DIR))
    else:
        log.info("Loading tokenizer %s (+ custom tokens)", MODEL_NAME)
        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
        tokenizer.add_tokens(CUSTOM_TOKENS)

    return dataset, class_weights, tokenizer


def build_model(tokenizer: AutoTokenizer) -> AutoModelForSequenceClassification:
    log.info("Loading model: %s", MODEL_NAME)
    model = AutoModelForSequenceClassification.from_pretrained(
        MODEL_NAME,
        num_labels=NUM_LABELS,
        id2label=ID2LABEL,
        label2id=LABEL2ID,
        ignore_mismatched_sizes=True,
    )

    orig = model.config.vocab_size
    model.resize_token_embeddings(len(tokenizer))
    new = model.get_input_embeddings().weight.shape[0]
    log.info("  Embeddings: %s -> %s tokens", orig, new)

    total = sum(p.numel() for p in model.parameters())
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    log.info("  Params: %.1fM total | %.1fM trainable", total / 1e6, trainable / 1e6)

    return model


def build_training_args() -> TrainingArguments:
    MODEL_OUT.mkdir(parents=True, exist_ok=True)
    LOGS_DIR.mkdir(parents=True, exist_ok=True)

    return TrainingArguments(
        output_dir=str(MODEL_OUT),
        num_train_epochs=CFG.num_epochs,
        per_device_train_batch_size=CFG.train_batch_size,
        per_device_eval_batch_size=CFG.eval_batch_size,
        learning_rate=CFG.learning_rate,
        optim="adamw_torch",
        warmup_ratio=CFG.warmup_ratio,
        weight_decay=CFG.weight_decay,
        max_grad_norm=CFG.max_grad_norm,
        evaluation_strategy="epoch",
        save_strategy="epoch",
        load_best_model_at_end=True,
        metric_for_best_model="f1",
        greater_is_better=True,
        save_total_limit=CFG.save_total_limit,
        logging_dir=str(LOGS_DIR),
        logging_steps=CFG.logging_steps,
        report_to="none",
        fp16=CFG.fp16,
        dataloader_num_workers=0,
        seed=CFG.seed,
        data_seed=CFG.seed,
    )


def full_evaluation(trainer: WeightedTrainer, dataset: DatasetDict) -> dict:
    log.info("Test set evaluation")
    test_results = trainer.evaluate(eval_dataset=dataset["test"])
    for k, v in sorted(test_results.items()):
        if not k.startswith("eval_"):
            continue
        log.info("  %s: %s", k, v)

    try:
        from sklearn.metrics import classification_report, confusion_matrix

        preds_output = trainer.predict(dataset["test"])
        preds = np.argmax(preds_output.predictions, axis=-1)
        labels = preds_output.label_ids

        print("\nClassification report:")
        print(
            classification_report(
                labels, preds, target_names=["safe", "scam"], digits=4
            )
        )

        cm = confusion_matrix(labels, preds)
        print("Confusion matrix (rows=actual, cols=predicted):")
        print(f"        safe  scam")
        print(f"  safe  {cm[0, 0]:5d} {cm[0, 1]:5d}")
        print(f"  scam  {cm[1, 0]:5d} {cm[1, 1]:5d}")

        denom = cm[1, 0] + cm[1, 1]
        fnr = cm[1, 0] / denom if denom else 0.0
        print(f"\nFalse negative rate (scam -> safe): {fnr * 100:.2f}%")

    except Exception as e:
        log.warning("Classification report skipped: %s", e)

    return test_results


def save_artifacts(model, tokenizer: AutoTokenizer, test_results: dict) -> None:
    log.info("Saving model -> %s", MODEL_OUT)
    trainer_eval_subset = {
        k: v
        for k, v in test_results.items()
        if k.startswith("eval_") and k not in ("eval_runtime", "eval_samples_per_second", "eval_steps_per_second")
    }
    model.save_pretrained(str(MODEL_OUT))
    tokenizer.save_pretrained(str(MODEL_OUT))

    metadata = {
        "model_name": MODEL_NAME,
        "num_labels": NUM_LABELS,
        "id2label": {str(k): v for k, v in ID2LABEL.items()},
        "label2id": LABEL2ID,
        "max_length": 128,
        "custom_tokens": CUSTOM_TOKENS,
        "train_config": {
            "num_epochs": CFG.num_epochs,
            "train_batch_size": CFG.train_batch_size,
            "eval_batch_size": CFG.eval_batch_size,
            "learning_rate": CFG.learning_rate,
            "warmup_ratio": CFG.warmup_ratio,
            "weight_decay": CFG.weight_decay,
            "fp16": CFG.fp16,
            "seed": CFG.seed,
        },
        "test_metrics": trainer_eval_subset,
    }

    meta_path = MODEL_OUT / "training_metadata.json"
    with open(meta_path, "w", encoding="utf-8") as f:
        json.dump(metadata, f, indent=2)
    log.info("Metadata -> %s", meta_path)


def smoke_test(model, tokenizer: AutoTokenizer) -> None:
    cases = [
        (
            "CBI officer here. You are under digital arrest for money laundering. "
            "Do not disconnect.",
            "scam",
        ),
        (
            "Your OTP is ready. Share it with me for KYC verification on your account.",
            "scam",
        ),
        ("Hey, are you coming to college tomorrow? Let me know.", "safe"),
        ("Your Amazon order has been shipped and will arrive by Friday.", "safe"),
    ]

    device = next(model.parameters()).device
    model.eval()

    print("\nSmoke test")
    print(f"{'Text':<58} exp   pred   scam%")
    print("-" * 88)

    for text, expected in cases:
        inputs = tokenizer(
            text,
            return_tensors="pt",
            truncation=True,
            padding="max_length",
            max_length=128,
        )
        inputs = {k: v.to(device) for k, v in inputs.items()}
        with torch.no_grad():
            logits = model(**inputs).logits
            probs = torch.softmax(logits, dim=-1)[0]
            pred_id = int(torch.argmax(probs).item())
            pred = ID2LABEL[pred_id]
            scam_p = probs[1].item()

        ok = "ok" if pred == expected else "XX"
        snippet = text[:56].replace("\n", " ")
        print(f"{snippet:<58} {expected:<5} {pred:<5} {scam_p * 100:5.1f}% {ok}")


def main() -> None:
    if hasattr(sys.stdout, "reconfigure"):
        try:
            sys.stdout.reconfigure(encoding="utf-8", errors="replace")
        except (OSError, ValueError):
            pass

    log.info("MuRIL scam classifier training")
    device = "cuda" if torch.cuda.is_available() else "cpu"
    log.info("Device: %s", device)
    if device == "cuda":
        log.info("GPU: %s", torch.cuda.get_device_name(0))

    dataset, class_weights, tokenizer = load_artifacts()
    model = build_model(tokenizer)

    training_args = build_training_args()
    log.info(
        "Config: epochs=%s batch=%s/%s lr=%s fp16=%s",
        CFG.num_epochs,
        CFG.train_batch_size,
        CFG.eval_batch_size,
        CFG.learning_rate,
        CFG.fp16,
    )

    trainer = WeightedTrainer(
        class_weights=class_weights,
        model=model,
        args=training_args,
        train_dataset=dataset["train"],
        eval_dataset=dataset["validation"],
        tokenizer=tokenizer,
        compute_metrics=build_compute_metrics(),
        callbacks=[
            EarlyStoppingCallback(early_stopping_patience=CFG.early_stopping_patience),
        ],
    )

    log.info("Training")
    train_result = trainer.train()
    log.info(
        "Done: steps=%s train_loss=%.4f time=%.1f min",
        train_result.global_step,
        train_result.training_loss,
        train_result.metrics.get("train_runtime", 0) / 60.0,
    )

    test_results = full_evaluation(trainer, dataset)
    save_artifacts(trainer.model, tokenizer, test_results)
    smoke_test(trainer.model, tokenizer)
    log.info("Finished.")


if __name__ == "__main__":
    main()