doc_intelligence/src/extraction/token_classifier.py

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
    }