Upload phase2/retrain_with_augmentation.py

phase2/retrain_with_augmentation.py

@@ -0,0 +1,208 @@
+"""
+Maskara Phase 2 Retraining with Targeted Augmentation
+=======================================================
+1. Downloads the current dataset from somukandula/maskara-indian-pii-200k.
+2. Generates targeted synthetic examples for under-performing entities.
+3. Merges them into the train split.
+4. Retrains from bert-base-uncased (or optionally from current Maskara weights).
+5. Evaluates and pushes the updated model to somukandula/maskara.
+
+Run on Modal GPU:
+  pip install torch transformers datasets evaluate trackio accelerate
+  python phase2/retrain_with_augmentation.py
+"""
+import os
+import json
+import warnings
+
+import numpy as np
+import evaluate
+from datasets import load_dataset, Dataset, concatenate_datasets
+from transformers import (
+    AutoTokenizer,
+    AutoModelForTokenClassification,
+    DataCollatorForTokenClassification,
+    TrainingArguments,
+    Trainer,
+    EarlyStoppingCallback,
+    set_seed,
+)
+
+warnings.filterwarnings("ignore", message=r".*unexpected keys in the checkpoint model loaded.*")
+warnings.filterwarnings("ignore", message=r".*missing keys in the checkpoint model loaded.*")
+
+MODEL_NAME = "google-bert/bert-base-uncased"  # change to "somukandula/maskara" for continued fine-tuning
+DATASET_NAME = "somukandula/maskara-indian-pii-200k"
+OUTPUT_DIR = "/app/maskara-phase2-retrain-output"
+HUB_MODEL_ID = "somukandula/maskara"
+TRACKIO_PROJECT = "maskara-phase2"
+
+SUPPORTED_ENTITIES = [
+    "ADDRESS", "API_KEY", "CREDIT_CARD", "DATE_OF_BIRTH", "DRIVER_LICENSE",
+    "EMAIL", "IP_ADDRESS", "PASSWORD", "PERSON_NAME", "PHONE", "SSN",
+    "USERNAME", "AADHAAR", "PAN_CARD", "PASSPORT", "UPI_ID", "VEHICLE_REG",
+]
+
+label_list = ["O"] + [f"B-{e}" for e in SUPPORTED_ENTITIES] + [f"I-{e}" for e in SUPPORTED_ENTITIES]
+label2id = {l: i for i, l in enumerate(label_list)}
+id2label = {i: l for l, i in label2id.items()}
+
+set_seed(2027)
+
+# Import augmentation generator
+import importlib.util
+spec = importlib.util.spec_from_file_location("targeted_augmentation", "phase2/targeted_augmentation.py")
+aug = importlib.util.module_from_spec(spec)
+spec.loader.exec_module(aug)
+
+
+def tokenize_and_align_labels(examples, tokenizer):
+    tokenized = tokenizer(
+        examples["text"],
+        truncation=True,
+        max_length=256,
+        return_offsets_mapping=True,
+    )
+    all_labels = []
+    for i in range(len(examples["text"])):
+        offsets = tokenized.offset_mapping[i]
+        labels = ["O"] * len(offsets)
+        for ent in examples["entities"][i]:
+            es, ee, el = ent["start"], ent["end"], ent["label"]
+            if el not in SUPPORTED_ENTITIES:
+                continue
+            first = True
+            for j, (ts, te) in enumerate(offsets):
+                if ts == te == 0:
+                    continue
+                if te <= es or ts >= ee:
+                    continue
+                labels[j] = f"B-{el}" if first else f"I-{el}"
+                first = False
+        all_labels.append([label2id.get(l, 0) for l in labels])
+    tokenized["labels"] = all_labels
+    tokenized.pop("offset_mapping", None)
+    return tokenized
+
+
+def main():
+    print("Loading tokenizer and dataset...")
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_fast=True)
+    dataset = load_dataset(DATASET_NAME)
+
+    print("Generating targeted augmentation...")
+    rng = __import__("random").Random(2027)
+    aug_examples = aug.generate_targeted_examples(rng, n_per_entity=2000)
+    aug_dataset = Dataset.from_list(aug_examples)
+
+    print(f"Original train size: {len(dataset['train'])}")
+    train = concatenate_datasets([dataset["train"], aug_dataset])
+    print(f"Augmented train size: {len(train)}")
+
+    tokenized = {}
+    for split in ["train", "template_disjoint_eval", "real_world_eval"]:
+        src = train if split == "train" else dataset[split]
+        tokenized[split] = src.map(
+            lambda ex: tokenize_and_align_labels(ex, tokenizer),
+            batched=True,
+            remove_columns=src.column_names,
+            num_proc=4,
+            desc=f"tokenize {split}",
+        )
+
+    seqeval = evaluate.load("seqeval")
+
+    def compute_metrics(p):
+        predictions, labels = p
+        predictions = np.argmax(predictions, axis=2)
+        true_predictions = [
+            [label_list[p] for p, lab in zip(prediction, label) if lab != -100]
+            for prediction, label in zip(predictions, labels)
+        ]
+        true_labels = [
+            [label_list[lab] for p, lab in zip(prediction, label) if lab != -100]
+            for prediction, label in zip(predictions, labels)
+        ]
+        results = seqeval.compute(predictions=true_predictions, references=true_labels, zero_division=0)
+        out = {
+            "precision": results["overall_precision"],
+            "recall": results["overall_recall"],
+            "f1": results["overall_f1"],
+            "accuracy": results["overall_accuracy"],
+        }
+        for key, val in results.items():
+            if isinstance(val, dict):
+                out[f"{key}_f1"] = val.get("f1", 0.0)
+        return out
+
+    print("Loading model...")
+    model = AutoModelForTokenClassification.from_pretrained(
+        MODEL_NAME,
+        num_labels=len(label_list),
+        id2label=id2label,
+        label2id=label2id,
+        ignore_mismatched_sizes=True,
+    )
+
+    training_args = TrainingArguments(
+        output_dir=OUTPUT_DIR,
+        overwrite_output_dir=True,
+        num_train_epochs=5,
+        per_device_train_batch_size=32,
+        per_device_eval_batch_size=64,
+        learning_rate=3e-5,
+        weight_decay=0.01,
+        warmup_ratio=0.1,
+        lr_scheduler_type="linear",
+        eval_strategy="epoch",
+        save_strategy="epoch",
+        logging_strategy="steps",
+        logging_steps=50,
+        logging_first_step=True,
+        disable_tqdm=True,
+        load_best_model_at_end=True,
+        metric_for_best_model="f1",
+        greater_is_better=True,
+        bf16=True,
+        optim="adamw_torch",
+        report_to="trackio",
+        run_name="maskara-phase2-retrain-targeted",
+        project=TRACKIO_PROJECT,
+        push_to_hub=True,
+        hub_model_id=HUB_MODEL_ID,
+        hub_strategy="end",
+        seed=2027,
+    )
+
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=tokenized["train"],
+        eval_dataset=tokenized["template_disjoint_eval"],
+        processing_class=tokenizer,
+        data_collator=DataCollatorForTokenClassification(tokenizer, pad_to_multiple_of=8),
+        compute_metrics=compute_metrics,
+        callbacks=[EarlyStoppingCallback(early_stopping_patience=2)],
+    )
+
+    print("Training...")
+    trainer.train()
+
+    td_metrics = trainer.evaluate(tokenized["template_disjoint_eval"])
+    rw_metrics = trainer.evaluate(tokenized["real_world_eval"])
+
+    metrics = {
+        "template_disjoint_eval": td_metrics,
+        "real_world_eval": rw_metrics,
+        "label_list": label_list,
+        "num_parameters": sum(p.numel() for p in model.parameters()),
+    }
+    with open(os.path.join(OUTPUT_DIR, "retrain_metrics.json"), "w") as f:
+        json.dump(metrics, f, indent=2, default=float)
+
+    trainer.push_to_hub(commit_message="Maskara Phase 2 retrain with targeted augmentation")
+    print("Done. Pushed to", HUB_MODEL_ID)
+
+
+if __name__ == "__main__":
+    main()