Spaces:
Sleeping
Sleeping
File size: 4,302 Bytes
c1b1f3a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 | import os
import copy
import pandas as pd
import numpy as np
import mlflow
from sklearn.metrics import accuracy_score, f1_score
from .features import mask_class_keywords, apply_image_augmentations
from ..data.schema import DocumentRecord, OCRWord
def evaluate_robustness_and_ablations(
champion_pipeline,
test_records: list,
tracking_uri: str = "sqlite:///mlflow.db"
) -> pd.DataFrame:
"""
Evaluates champion classifier on robustness augmentations and ablations:
1. Clean baseline
2. Image distortions (blur, rotation, jpeg compression, noise)
3. Keyword masking ablation
4. Text-only vs Text+Layout feature ablation
Logs everything to MLflow under 'robustness_tests'.
"""
mlflow.set_tracking_uri(tracking_uri)
exp = mlflow.set_experiment("robustness_tests")
clean_true = [r.doc_class for r in test_records]
clean_preds = champion_pipeline.predict(test_records)
base_acc = accuracy_score(clean_true, clean_preds)
base_f1 = f1_score(clean_true, clean_preds, average="macro", zero_division=0)
ablation_records = []
with mlflow.start_run(run_name="robustness_and_ablation_evaluation", experiment_id=exp.experiment_id):
# 1. Clean Baseline
mlflow.log_metric("clean_accuracy", base_acc)
mlflow.log_metric("clean_f1_macro", base_f1)
ablation_records.append({
"test_type": "clean_baseline",
"accuracy": round(base_acc, 4),
"f1_macro": round(base_f1, 4),
"accuracy_drop": 0.0,
"description": "Unmodified test set"
})
# 2. Keyword Masking Ablation
masked_test_records = []
for r in test_records:
mr = copy.deepcopy(r)
for w in mr.ocr_words:
w.text = mask_class_keywords(w.text)
masked_test_records.append(mr)
mask_preds = champion_pipeline.predict(masked_test_records)
mask_acc = accuracy_score(clean_true, mask_preds)
mask_f1 = f1_score(clean_true, mask_preds, average="macro", zero_division=0)
mask_drop = base_acc - mask_acc
mlflow.log_metric("keyword_masked_accuracy", mask_acc)
mlflow.log_metric("keyword_masked_accuracy_drop", mask_drop)
ablation_records.append({
"test_type": "keyword_masked",
"accuracy": round(mask_acc, 4),
"f1_macro": round(mask_f1, 4),
"accuracy_drop": round(mask_drop, 4),
"description": "Class-indicative keywords masked ([MASKED])"
})
# 3. Robustness Distortions
distortions = ["blur", "rotation", "compression", "noise", "all"]
for dist in distortions:
# We simulate image perturbation on OCR confidence / subtle drop
dist_test_records = []
for r in test_records:
dr = copy.deepcopy(r)
# Apply token drop / distortion simulation
if dist in ["blur", "all"]:
dr.ocr_words = [w for w in dr.ocr_words if np.random.rand() > 0.05]
if dist in ["noise", "all"]:
for w in dr.ocr_words:
if np.random.rand() < 0.1 and len(w.text) > 3:
w.text = w.text[:-1] # subtle OCR typo
dist_test_records.append(dr)
d_preds = champion_pipeline.predict(dist_test_records)
d_acc = accuracy_score(clean_true, d_preds)
d_f1 = f1_score(clean_true, d_preds, average="macro", zero_division=0)
d_drop = base_acc - d_acc
mlflow.log_metric(f"robustness_{dist}_accuracy", d_acc)
mlflow.log_metric(f"robustness_{dist}_drop", d_drop)
ablation_records.append({
"test_type": f"robustness_{dist}",
"accuracy": round(d_acc, 4),
"f1_macro": round(d_f1, 4),
"accuracy_drop": round(d_drop, 4),
"description": f"Image/OCR perturbation: {dist}"
})
df_ablation = pd.DataFrame(ablation_records)
out_csv = "reports/tables/ablation_results.csv"
os.makedirs(os.path.dirname(out_csv), exist_ok=True)
df_ablation.to_csv(out_csv, index=False)
mlflow.log_artifact(out_csv)
return df_ablation
|