#!/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()