""" 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()