addyo07/query-scope-classifier

1#!/usr/bin/env python32"""3Layer 2: Baseline Evaluation, GPU Fine-Tuning & Dynamics Optimization for ModernBERT-base4Master Golden Dataset: /opt/vox/sandbox/datasets/memory_scope_golden_v1.json (22,006 samples)5"""6 7import os8import sys9import json10import time11import torch12import numpy as np13import pandas as pd14from datasets import Dataset15from transformers import (16    AutoTokenizer,17    AutoModelForSequenceClassification,18    Trainer,19    TrainingArguments,20    DataCollatorWithPadding,21)22from sklearn.metrics import accuracy_score, precision_recall_fscore_support, classification_report23from sklearn.model_selection import train_test_split24 25GOLDEN_DATASET_PATH = "/opt/vox/sandbox/datasets/memory_scope_golden_v1.json"26BASE_MODEL_NAME = "answerdotai/ModernBERT-base"27OUTPUT_DIR = "/opt/vox/sandbox/artifacts/modernbert_scope_final"28RESULTS_DIR = "/opt/vox/sandbox/results"29 30os.makedirs(OUTPUT_DIR, exist_ok=True)31os.makedirs(RESULTS_DIR, exist_ok=True)32 33SCOPE_MAP = {"ChitChat": 0, "User": 1, "Domain": 2, "Temporal": 3}34ID_TO_SCOPE = {0: "ChitChat", 1: "User", 2: "Domain", 3: "Temporal"}35 36def compute_metrics(eval_pred):37    logits, labels = eval_pred38    preds = np.argmax(logits, axis=1)39    40    precision, recall, f1, _ = precision_recall_fscore_support(41        labels, preds, average="macro", zero_division=042    )43    acc = accuracy_score(labels, preds)44    45    _, class_recall, _, _ = precision_recall_fscore_support(46        labels, preds, average=None, labels=[0, 1, 2, 3], zero_division=047    )48    49    return {50        "accuracy": acc,51        "macro_f1": f1,52        "macro_precision": precision,53        "macro_recall": recall,54        "recall_chitchat": class_recall[0],55        "recall_user": class_recall[1],56        "recall_domain": class_recall[2],57        "recall_temporal": class_recall[3],58    }59 60def main():61    print("=== Layer 2: Baseline Evaluation & GPU Fine-Tuning Pipeline (ModernBERT-base) ===", flush=True)62    63    # 1. Load Master Golden Dataset64    if not os.path.exists(GOLDEN_DATASET_PATH):65        print(f"Error: {GOLDEN_DATASET_PATH} missing!", flush=True)66        sys.exit(1)67        68    with open(GOLDEN_DATASET_PATH, "r", encoding="utf-8") as f:69        data_payload = json.load(f)70        samples = data_payload["samples"]71        72    print(f"Loaded {len(samples)} total samples from Master Golden Dataset.", flush=True)73    74    formatted_data = [75        {76            "id": s["id"],77            "text": s["text"],78            "label": SCOPE_MAP[s["scope"]],79            "language": s.get("language", "en"),80            "strat_key": f"{s['scope']}_{s.get('language', 'en')}"81        }82        for s in samples83    ]84    85    df = pd.DataFrame(formatted_data)86    87    # 80% Train (17,604), 10% Val (2,201), 10% Test (2,201)88    train_df, temp_df = train_test_split(df, test_size=0.20, random_state=42, stratify=df["strat_key"])89    val_df, test_df = train_test_split(temp_df, test_size=0.50, random_state=42, stratify=temp_df["strat_key"])90    91    print(f"Dataset Split: Train={len(train_df)}, Val={len(val_df)}, Test={len(test_df)}", flush=True)92    93    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_NAME)94    95    def tokenize_df(df_input):96        ds = Dataset.from_pandas(df_input)97        ds_mapped = ds.map(98            lambda x: tokenizer(x["text"], truncation=True, max_length=64, padding=False),99            batched=True,100        )101        cols_to_keep = ["input_ids", "attention_mask", "label"]102        cols_to_remove = [c for c in ds_mapped.column_names if c not in cols_to_keep]103        return ds_mapped.remove_columns(cols_to_remove)104        105    train_ds = tokenize_df(train_df)106    val_ds = tokenize_df(val_df)107    test_ds = tokenize_df(test_df)108    109    # 2. Phase 2.1: Pretrained Zero-Shot Baseline Evaluation110    print("\n--- Phase 2.1: Zero-Shot Baseline Evaluation of Pretrained ModernBERT-base ---", flush=True)111    baseline_model = AutoModelForSequenceClassification.from_pretrained(112        BASE_MODEL_NAME,113        num_labels=4,114        id2label=ID_TO_SCOPE,115        label2id=SCOPE_MAP,116    )117    118    trainer_baseline = Trainer(119        model=baseline_model,120        processing_class=tokenizer,121        data_collator=DataCollatorWithPadding(tokenizer=tokenizer),122        compute_metrics=compute_metrics,123    )124    125    baseline_eval = trainer_baseline.evaluate(test_ds)126    print("Baseline Zero-Shot Test Evaluation Results:")127    for k, v in baseline_eval.items():128        print(f"  - {k}: {v}", flush=True)129        130    with open(os.path.join(RESULTS_DIR, "baseline_zero_shot_eval.json"), "w") as f:131        json.dump(baseline_eval, f, indent=2)132        133    # 3. Phase 2.2: GPU Fine-Tuning Execution on RTX 5070 Ti134    print("\n--- Phase 2.2: GPU Fine-Tuning Execution on RTX 5070 Ti ---", flush=True)135    136    model = AutoModelForSequenceClassification.from_pretrained(137        BASE_MODEL_NAME,138        num_labels=4,139        id2label=ID_TO_SCOPE,140        label2id=SCOPE_MAP,141    )142    143    training_args = TrainingArguments(144        output_dir=OUTPUT_DIR,145        eval_strategy="epoch",146        save_strategy="no",147        learning_rate=3e-5,148        per_device_train_batch_size=32,149        per_device_eval_batch_size=64,150        num_train_epochs=3,151        weight_decay=0.01,152        warmup_ratio=0.10,153        logging_steps=50,154        bf16=True,155        report_to="none",156    )157    158    trainer = Trainer(159        model=model,160        args=training_args,161        train_dataset=train_ds,162        eval_dataset=val_ds,163        processing_class=tokenizer,164        data_collator=DataCollatorWithPadding(tokenizer=tokenizer),165        compute_metrics=compute_metrics,166    )167    168    print("Starting fine-tuning training loop...", flush=True)169    trainer.train()170    171    final_model_path = os.path.join(OUTPUT_DIR, "final_pytorch_model")172    trainer.save_model(final_model_path)173    tokenizer.save_pretrained(final_model_path)174    print(f"Fine-tuned PyTorch model saved to {final_model_path}", flush=True)175    176    # 4. Phase 2.3: Holdout Test Set Evaluation & Gate Audit177    print("\n--- Phase 2.3: Fine-Tuned Holdout Test Evaluation & Gate 2 Audit ---", flush=True)178    final_eval = trainer.evaluate(test_ds)179    180    print("\nFinal Fine-Tuned Test Metrics:")181    for k, v in final_eval.items():182        print(f"  - {k}: {v}", flush=True)183        184    with open(os.path.join(RESULTS_DIR, "finetuned_test_eval.json"), "w") as f:185        json.dump(final_eval, f, indent=2)186        187    test_acc = final_eval.get("eval_accuracy", 0.0)188    test_f1 = final_eval.get("eval_macro_f1", 0.0)189    190    print("\n==================================================================", flush=True)191    print(f"🎯 LAYER 2 MILESTONE VERDICT: {'✅ PASSED' if (test_acc >= 0.88 and test_f1 >= 0.88) else '❌ FAILED'}", flush=True)192    print(f"  - Holdout Test Accuracy: {test_acc*100:.2f}% (Target: ≥88.0%)", flush=True)193    print(f"  - Holdout Macro F1: {test_f1:.4f} (Target: ≥0.8800)", flush=True)194    print(f"  - Baseline Net Gain: Accuracy +{(test_acc - baseline_eval.get('eval_accuracy', 0.0))*100:.2f}%", flush=True)195    print("==================================================================", flush=True)196 197if __name__ == "__main__":198    main()199