Spaces:
Sleeping
Sleeping
| 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 | |