scripts/phase2_train_modernbert_scope.py

#!/usr/bin/env python3
"""
Layer 2: Baseline Evaluation, GPU Fine-Tuning & Dynamics Optimization for ModernBERT-base
Master Golden Dataset: /opt/vox/sandbox/datasets/memory_scope_golden_v1.json (22,006 samples)
"""

import os
import sys
import json
import time
import torch
import numpy as np
import pandas as pd
from datasets import Dataset
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    Trainer,
    TrainingArguments,
    DataCollatorWithPadding,
)
from sklearn.metrics import accuracy_score, precision_recall_fscore_support, classification_report
from sklearn.model_selection import train_test_split

GOLDEN_DATASET_PATH = "/opt/vox/sandbox/datasets/memory_scope_golden_v1.json"
BASE_MODEL_NAME = "answerdotai/ModernBERT-base"
OUTPUT_DIR = "/opt/vox/sandbox/artifacts/modernbert_scope_final"
RESULTS_DIR = "/opt/vox/sandbox/results"

os.makedirs(OUTPUT_DIR, exist_ok=True)
os.makedirs(RESULTS_DIR, exist_ok=True)

SCOPE_MAP = {"ChitChat": 0, "User": 1, "Domain": 2, "Temporal": 3}
ID_TO_SCOPE = {0: "ChitChat", 1: "User", 2: "Domain", 3: "Temporal"}

def compute_metrics(eval_pred):
    logits, labels = eval_pred
    preds = np.argmax(logits, axis=1)
    
    precision, recall, f1, _ = precision_recall_fscore_support(
        labels, preds, average="macro", zero_division=0
    )
    acc = accuracy_score(labels, preds)
    
    _, class_recall, _, _ = precision_recall_fscore_support(
        labels, preds, average=None, labels=[0, 1, 2, 3], zero_division=0
    )
    
    return {
        "accuracy": acc,
        "macro_f1": f1,
        "macro_precision": precision,
        "macro_recall": recall,
        "recall_chitchat": class_recall[0],
        "recall_user": class_recall[1],
        "recall_domain": class_recall[2],
        "recall_temporal": class_recall[3],
    }

def main():
    print("=== Layer 2: Baseline Evaluation & GPU Fine-Tuning Pipeline (ModernBERT-base) ===", flush=True)
    
    # 1. Load Master Golden Dataset
    if not os.path.exists(GOLDEN_DATASET_PATH):
        print(f"Error: {GOLDEN_DATASET_PATH} missing!", flush=True)
        sys.exit(1)
        
    with open(GOLDEN_DATASET_PATH, "r", encoding="utf-8") as f:
        data_payload = json.load(f)
        samples = data_payload["samples"]
        
    print(f"Loaded {len(samples)} total samples from Master Golden Dataset.", flush=True)
    
    formatted_data = [
        {
            "id": s["id"],
            "text": s["text"],
            "label": SCOPE_MAP[s["scope"]],
            "language": s.get("language", "en"),
            "strat_key": f"{s['scope']}_{s.get('language', 'en')}"
        }
        for s in samples
    ]
    
    df = pd.DataFrame(formatted_data)
    
    # 80% Train (17,604), 10% Val (2,201), 10% Test (2,201)
    train_df, temp_df = train_test_split(df, test_size=0.20, random_state=42, stratify=df["strat_key"])
    val_df, test_df = train_test_split(temp_df, test_size=0.50, random_state=42, stratify=temp_df["strat_key"])
    
    print(f"Dataset Split: Train={len(train_df)}, Val={len(val_df)}, Test={len(test_df)}", flush=True)
    
    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_NAME)
    
    def tokenize_df(df_input):
        ds = Dataset.from_pandas(df_input)
        ds_mapped = ds.map(
            lambda x: tokenizer(x["text"], truncation=True, max_length=64, padding=False),
            batched=True,
        )
        cols_to_keep = ["input_ids", "attention_mask", "label"]
        cols_to_remove = [c for c in ds_mapped.column_names if c not in cols_to_keep]
        return ds_mapped.remove_columns(cols_to_remove)
        
    train_ds = tokenize_df(train_df)
    val_ds = tokenize_df(val_df)
    test_ds = tokenize_df(test_df)
    
    # 2. Phase 2.1: Pretrained Zero-Shot Baseline Evaluation
    print("\n--- Phase 2.1: Zero-Shot Baseline Evaluation of Pretrained ModernBERT-base ---", flush=True)
    baseline_model = AutoModelForSequenceClassification.from_pretrained(
        BASE_MODEL_NAME,
        num_labels=4,
        id2label=ID_TO_SCOPE,
        label2id=SCOPE_MAP,
    )
    
    trainer_baseline = Trainer(
        model=baseline_model,
        processing_class=tokenizer,
        data_collator=DataCollatorWithPadding(tokenizer=tokenizer),
        compute_metrics=compute_metrics,
    )
    
    baseline_eval = trainer_baseline.evaluate(test_ds)
    print("Baseline Zero-Shot Test Evaluation Results:")
    for k, v in baseline_eval.items():
        print(f"  - {k}: {v}", flush=True)
        
    with open(os.path.join(RESULTS_DIR, "baseline_zero_shot_eval.json"), "w") as f:
        json.dump(baseline_eval, f, indent=2)
        
    # 3. Phase 2.2: GPU Fine-Tuning Execution on RTX 5070 Ti
    print("\n--- Phase 2.2: GPU Fine-Tuning Execution on RTX 5070 Ti ---", flush=True)
    
    model = AutoModelForSequenceClassification.from_pretrained(
        BASE_MODEL_NAME,
        num_labels=4,
        id2label=ID_TO_SCOPE,
        label2id=SCOPE_MAP,
    )
    
    training_args = TrainingArguments(
        output_dir=OUTPUT_DIR,
        eval_strategy="epoch",
        save_strategy="no",
        learning_rate=3e-5,
        per_device_train_batch_size=32,
        per_device_eval_batch_size=64,
        num_train_epochs=3,
        weight_decay=0.01,
        warmup_ratio=0.10,
        logging_steps=50,
        bf16=True,
        report_to="none",
    )
    
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_ds,
        eval_dataset=val_ds,
        processing_class=tokenizer,
        data_collator=DataCollatorWithPadding(tokenizer=tokenizer),
        compute_metrics=compute_metrics,
    )
    
    print("Starting fine-tuning training loop...", flush=True)
    trainer.train()
    
    final_model_path = os.path.join(OUTPUT_DIR, "final_pytorch_model")
    trainer.save_model(final_model_path)
    tokenizer.save_pretrained(final_model_path)
    print(f"Fine-tuned PyTorch model saved to {final_model_path}", flush=True)
    
    # 4. Phase 2.3: Holdout Test Set Evaluation & Gate Audit
    print("\n--- Phase 2.3: Fine-Tuned Holdout Test Evaluation & Gate 2 Audit ---", flush=True)
    final_eval = trainer.evaluate(test_ds)
    
    print("\nFinal Fine-Tuned Test Metrics:")
    for k, v in final_eval.items():
        print(f"  - {k}: {v}", flush=True)
        
    with open(os.path.join(RESULTS_DIR, "finetuned_test_eval.json"), "w") as f:
        json.dump(final_eval, f, indent=2)
        
    test_acc = final_eval.get("eval_accuracy", 0.0)
    test_f1 = final_eval.get("eval_macro_f1", 0.0)
    
    print("\n==================================================================", flush=True)
    print(f"🎯 LAYER 2 MILESTONE VERDICT: {'✅ PASSED' if (test_acc >= 0.88 and test_f1 >= 0.88) else '❌ FAILED'}", flush=True)
    print(f"  - Holdout Test Accuracy: {test_acc*100:.2f}% (Target: ≥88.0%)", flush=True)
    print(f"  - Holdout Macro F1: {test_f1:.4f} (Target: ≥0.8800)", flush=True)
    print(f"  - Baseline Net Gain: Accuracy +{(test_acc - baseline_eval.get('eval_accuracy', 0.0))*100:.2f}%", flush=True)
    print("==================================================================", flush=True)

if __name__ == "__main__":
    main()