Spaces:
Sleeping
Sleeping
| import os | |
| import json | |
| import time | |
| import re | |
| import numpy as np | |
| import pandas as pd | |
| import matplotlib.pyplot as plt | |
| import cv2 | |
| from typing import List, Dict, Any, Tuple | |
| import mlflow | |
| from sklearn.ensemble import RandomForestClassifier | |
| from sklearn.feature_extraction import DictVectorizer | |
| from sklearn.metrics import classification_report, f1_score | |
| from .field_metrics import evaluate_field_extraction, normalized_levenshtein_similarity | |
| from .postprocess import postprocess_extracted_fields | |
| from .rule_based import RuleBasedExtractor | |
| from .proximity_extractor import ProximityHeuristicsExtractor | |
| from ..data.schema import DocumentRecord, OCRWord, CLASS_FIELD_SCHEMAS | |
| def fuzzy_align_token(token: str, target: str) -> bool: | |
| """Checks if token is part of the ground truth target field.""" | |
| if not target or not token: | |
| return False | |
| t_clean = re.sub(r'[^\w]', '', token.lower()) | |
| target_clean = re.sub(r'[^\w]', '', str(target).lower()) | |
| if len(t_clean) >= 2 and t_clean in target_clean: | |
| return True | |
| return False | |
| def extract_word_features(w: OCRWord, prev_w: OCRWord = None, next_w: OCRWord = None) -> Dict[str, Any]: | |
| text = w.text | |
| x0, y0, x1, y1 = w.bbox | |
| w_box = max(x1 - x0, 1) | |
| h_box = max(y1 - y0, 1) | |
| feats = { | |
| "text_lower": text.lower()[:8], | |
| "is_digit": text.isdigit(), | |
| "is_title": text.istitle(), | |
| "is_upper": text.isupper(), | |
| "has_curr": any(c in text for c in "$€£¥"), | |
| "has_slash": "/" in text or "-" in text, | |
| "len": len(text), | |
| "norm_x0": round(x0 / 1000.0, 3), | |
| "norm_y0": round(y0 / 1000.0, 3), | |
| "norm_x1": round(x1 / 1000.0, 3), | |
| "norm_y1": round(y1 / 1000.0, 3), | |
| "aspect_ratio": round(w_box / h_box, 2) | |
| } | |
| if prev_w: | |
| feats["prev_lower"] = prev_w.text.lower()[:6] | |
| feats["prev_is_digit"] = prev_w.text.isdigit() | |
| else: | |
| feats["prev_lower"] = "<START>" | |
| feats["prev_is_digit"] = False | |
| if next_w: | |
| feats["next_lower"] = next_w.text.lower()[:6] | |
| else: | |
| feats["next_lower"] = "<END>" | |
| return feats | |
| class MLTokenClassifierExtractor: | |
| """ | |
| ML Token Classifier: Extracts word features and trains on BIO tags | |
| to predict structured document fields. | |
| """ | |
| def __init__(self, doc_class: str): | |
| self.doc_class = doc_class | |
| self.schema_fields = CLASS_FIELD_SCHEMAS.get(doc_class, []) | |
| self.vectorizer = DictVectorizer(sparse=False) | |
| self.clf = RandomForestClassifier(n_estimators=100, max_depth=12, random_state=42, class_weight="balanced") | |
| self.tag_to_idx = {"O": 0} | |
| for f in self.schema_fields: | |
| self.tag_to_idx[f"B-{f}"] = len(self.tag_to_idx) | |
| self.tag_to_idx[f"I-{f}"] = len(self.tag_to_idx) | |
| self.idx_to_tag = {v: k for k, v in self.tag_to_idx.items()} | |
| def _align_bio_tags(self, record: DocumentRecord) -> Tuple[List[str], float]: | |
| tags = ["O"] * len(record.ocr_words) | |
| aligned_count = 0 | |
| total_fields = sum(1 for v in record.fields.values() if v is not None) | |
| for field_name, field_val in record.fields.items(): | |
| if not field_val or field_name not in self.schema_fields: | |
| continue | |
| first_match = True | |
| matched = False | |
| for i, w in enumerate(record.ocr_words): | |
| if fuzzy_align_token(w.text, str(field_val)): | |
| tag = f"B-{field_name}" if first_match else f"I-{field_name}" | |
| tags[i] = tag | |
| first_match = False | |
| matched = True | |
| if matched: | |
| aligned_count += 1 | |
| coverage = aligned_count / max(total_fields, 1) | |
| return tags, coverage | |
| def fit(self, records: List[DocumentRecord]) -> float: | |
| class_records = [r for r in records if r.doc_class == self.doc_class] | |
| if not class_records: | |
| return 0.0 | |
| all_feats = [] | |
| all_labels = [] | |
| coverages = [] | |
| for r in class_records: | |
| tags, cov = self._align_bio_tags(r) | |
| coverages.append(cov) | |
| words = r.ocr_words | |
| for i, w in enumerate(words): | |
| prev_w = words[i - 1] if i > 0 else None | |
| next_w = words[i + 1] if i < len(words) - 1 else None | |
| feat = extract_word_features(w, prev_w, next_w) | |
| all_feats.append(feat) | |
| all_labels.append(self.tag_to_idx.get(tags[i], 0)) | |
| X = self.vectorizer.fit_transform(all_feats) | |
| y = np.array(all_labels) | |
| self.clf.fit(X, y) | |
| return float(np.mean(coverages)) | |
| def extract(self, record: DocumentRecord) -> Dict[str, Any]: | |
| words = record.ocr_words | |
| if not words or not hasattr(self.clf, "classes_"): | |
| return {f: None for f in self.schema_fields} | |
| feats = [] | |
| for i, w in enumerate(words): | |
| prev_w = words[i - 1] if i > 0 else None | |
| next_w = words[i + 1] if i < len(words) - 1 else None | |
| feats.append(extract_word_features(w, prev_w, next_w)) | |
| X = self.vectorizer.transform(feats) | |
| preds = self.clf.predict(X) | |
| extracted = {f: [] for f in self.schema_fields} | |
| for i, p_idx in enumerate(preds): | |
| tag = self.idx_to_tag.get(p_idx, "O") | |
| if tag != "O": | |
| f_name = tag.split("-")[1] | |
| if f_name in extracted: | |
| extracted[f_name].append(words[i].text) | |
| return {k: " ".join(v) if v else None for k, v in extracted.items()} | |
| def run_extraction_experiments( | |
| train_records: List[DocumentRecord], | |
| test_records: List[DocumentRecord], | |
| artifacts_dir: str = "reports/extraction_artifacts", | |
| tracking_uri: str = "sqlite:///mlflow.db" | |
| ) -> Dict[str, Any]: | |
| """ | |
| Runs comparison between Rule-Based (Baseline A), Proximity (Baseline B), | |
| and ML Token Classifier on held-out test records. | |
| """ | |
| os.makedirs(artifacts_dir, exist_ok=True) | |
| mlflow.set_tracking_uri(tracking_uri) | |
| exp = mlflow.set_experiment("field_extraction") | |
| rule_extractor = RuleBasedExtractor() | |
| prox_extractor = ProximityHeuristicsExtractor() | |
| classes = ["invoice", "receipt", "application_form"] | |
| ml_extractors = {} | |
| alignment_coverages = {} | |
| for c in classes: | |
| ext = MLTokenClassifierExtractor(c) | |
| cov = ext.fit(train_records) | |
| ml_extractors[c] = ext | |
| alignment_coverages[c] = cov | |
| models = ["rule_based", "proximity_heuristics", "ml_token_classifier"] | |
| summary_results = [] | |
| error_analysis_rows = [] | |
| for mname in models: | |
| with mlflow.start_run(run_name=f"ext_{mname}", experiment_id=exp.experiment_id): | |
| mlflow.set_tag("model_family", mname) | |
| mlflow.set_tag("stage", "stage3_field_extraction") | |
| all_preds = [] | |
| all_golds = [] | |
| all_classes = [] | |
| t0 = time.time() | |
| for r in test_records: | |
| cls_name = r.doc_class | |
| if mname == "rule_based": | |
| raw_pred = rule_extractor.extract(r, cls_name) | |
| elif mname == "proximity_heuristics": | |
| raw_pred = prox_extractor.extract(r, cls_name) | |
| else: | |
| raw_pred = ml_extractors[cls_name].extract(r) | |
| # Post-process | |
| norm_pred, val_res = postprocess_extracted_fields(cls_name, raw_pred) | |
| all_preds.append(norm_pred) | |
| all_golds.append(r.fields) | |
| all_classes.append(cls_name) | |
| # Log error analysis sample | |
| for f_name, gold_val in r.fields.items(): | |
| if gold_val is not None: | |
| pred_val = norm_pred.get(f_name) | |
| lev = normalized_levenshtein_similarity(str(pred_val or ""), str(gold_val)) | |
| if lev < 0.8: | |
| error_analysis_rows.append({ | |
| "doc_id": r.doc_id, | |
| "class": cls_name, | |
| "field": f_name, | |
| "gold": gold_val, | |
| "predicted": pred_val, | |
| "similarity": round(lev, 3), | |
| "model": mname | |
| }) | |
| latency_ms = (time.time() - t0) * 1000.0 / max(len(test_records), 1) | |
| # Evaluate per class | |
| class_evals = {} | |
| for c in classes: | |
| c_indices = [i for i, x in enumerate(all_classes) if x == c] | |
| c_preds = [all_preds[i] for i in c_indices] | |
| c_golds = [all_golds[i] for i in c_indices] | |
| res = evaluate_field_extraction(c_preds, c_golds, CLASS_FIELD_SCHEMAS[c]) | |
| class_evals[c] = res | |
| mlflow.log_metric(f"{c}_macro_token_f1", res["macro_token_f1"]) | |
| mlflow.log_metric(f"{c}_macro_exact_match", res["macro_exact_match"]) | |
| overall_f1 = float(np.mean([class_evals[c]["macro_token_f1"] for c in classes])) | |
| overall_em = float(np.mean([class_evals[c]["macro_exact_match"] for c in classes])) | |
| mlflow.log_metric("overall_macro_token_f1", overall_f1) | |
| mlflow.log_metric("overall_macro_exact_match", overall_em) | |
| mlflow.log_metric("inference_latency_ms_per_doc", latency_ms) | |
| summary_results.append({ | |
| "model": mname, | |
| "overall_macro_token_f1": round(overall_f1, 4), | |
| "overall_exact_match": round(overall_em, 4), | |
| "invoice_f1": round(class_evals["invoice"]["macro_token_f1"], 4), | |
| "receipt_f1": round(class_evals["receipt"]["macro_token_f1"], 4), | |
| "form_f1": round(class_evals["application_form"]["macro_token_f1"], 4), | |
| "latency_ms": round(latency_ms, 2) | |
| }) | |
| # Save summary tables and error analysis | |
| df_sum = pd.DataFrame(summary_results) | |
| sum_csv = "reports/tables/extraction_model_comparison.csv" | |
| os.makedirs(os.path.dirname(sum_csv), exist_ok=True) | |
| df_sum.to_csv(sum_csv, index=False) | |
| df_sum.to_markdown("reports/tables/extraction_model_comparison.md", index=False) | |
| df_err = pd.DataFrame(error_analysis_rows[:100]) | |
| os.makedirs(artifacts_dir, exist_ok=True) | |
| err_csv = os.path.join(artifacts_dir, "error_analysis.csv") | |
| df_err.to_csv(err_csv, index=False) | |
| # Field F1 Bar Chart | |
| plt.figure(figsize=(8, 4.5)) | |
| bars = plt.bar(df_sum["model"], df_sum["overall_macro_token_f1"], color=['#4C72B0', '#55A868', '#C44E52']) | |
| plt.ylabel("Macro Token F1") | |
| plt.title("Field Extraction Model Family Comparison") | |
| plt.ylim(0, 1.0) | |
| for bar in bars: | |
| yval = bar.get_height() | |
| plt.text(bar.get_x() + bar.get_width()/2.0, yval + 0.02, f"{yval:.3f}", ha='center', va='bottom') | |
| plt.tight_layout() | |
| chart_path = os.path.join(artifacts_dir, "field_f1_bar_chart.png") | |
| plt.savefig(chart_path, dpi=150) | |
| plt.close() | |
| return { | |
| "ml_extractors": ml_extractors, | |
| "summary": df_sum | |
| } | |