chore: capture semantic pipeline baseline

2026-09-21 15:52:17 -03:00
parent 9dcc263217
commit 6731eeca08
47 changed files with 3724 additions and 15 deletions
@@ -0,0 +1,115 @@
+"""NCM retrieval evaluation V2 (sections 1, 13).
+
+Same sampling as V1 (seed 42, per-segment) with each sample now quality-classified:
+  GOLD   - GTIN product or cross-tenant agreement on the same fingerprint/NCM
+  SILVER - single-tenant valid NCM, no independent confirmation
+  NOISY  - commercial description with < 2 informative tokens (not pursued)
+Retrieval measured separately: OFFICIAL_ONLY (hierarchical enriched corpus),
+KB_ONLY (neighbor bridge), HYBRID; recall@1/5/10/20/50 + MRR, per class and segment.
+"""
+import argparse
+import json
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
+
+from sqlalchemy import select
+
+from central.app.ai.retrieval import (EnrichedNcmCorpus, KbNeighborIndex, commercial_tokens, derive_abbreviation_lexicon,
+                                      hierarchical_candidates, hybrid_candidates, kb_neighbor_candidates)
+from central.app.db import SessionLocal
+from central.app.models import Dataset, KbObservation, KbProduct, NcmEntry
+
+SEGMENTS = {"CONSTRUCAO": "construcao", "MMMODAS": "moda", "RINCAO": "restaurante", "UNIAO": "mercado", "TESTESOFICIAIS": "base_testes"}
+KS = (1, 5, 10, 20, 50)
+
+
+def _samples(session, per_segment=40, seed=42):
+    import random
+    dataset = session.scalar(select(Dataset).where(Dataset.kind == "NCM", Dataset.status == "ACTIVE"))
+    valid_ncm = set(session.scalars(select(NcmEntry.code).where(NcmEntry.dataset_id == dataset.id)).all())
+    rows = session.execute(
+        select(KbProduct, KbObservation).join(KbObservation, KbObservation.kb_product_id == KbProduct.id)
+    ).all()
+    per_product_tenants = {}
+    for product, observation in rows:
+        per_product_tenants.setdefault(product.id, set()).add(observation.tenant_id)
+    by_segment = {}
+    for product, observation in rows:
+        ncm = (observation.observed_fields or {}).get("ncm")
+        if ncm and ncm in valid_ncm and product.normalized_description:
+            by_segment.setdefault(SEGMENTS.get(observation.database_label, observation.database_label), set()).add(
+                (product.id, product.normalized_description, ncm, bool(product.gtin), len(per_product_tenants[product.id]) >= 2)
+            )
+    rng = random.Random(seed)
+    samples = []
+    for segment, values in sorted(by_segment.items()):
+        picked = sorted(values) if len(values) <= per_segment else rng.sample(sorted(values), per_segment)
+        for product_id, description, ncm, has_gtin, cross_tenant in picked:
+            quality = "NOISY" if len(commercial_tokens(description)) < 2 else ("GOLD" if (has_gtin or cross_tenant) else "SILVER")
+            samples.append({"segment": segment, "product_id": product_id, "description": description, "true_ncm": ncm, "quality": quality})
+    return samples
+
+
+def _recall_mrr(runs, samples):
+    """runs: {mode: [ranked_codes_per_sample]}"""
+    out = {}
+    for mode, ranked_lists in runs.items():
+        metrics = {f"recall_at_{k}": 0 for k in KS} | {"mrr": 0.0}
+        total = 0
+        for sample, ranked in zip(samples, ranked_lists):
+            codes = [row[0] for row in ranked]
+            total += 1
+            for k in KS:
+                if sample["true_ncm"] in codes[:k]:
+                    metrics[f"recall_at_{k}"] += 1
+            if sample["true_ncm"] in codes:
+                metrics["mrr"] += 1.0 / (codes.index(sample["true_ncm"]) + 1)
+        if total:
+            for key in metrics:
+                metrics[key] = round(metrics[key] / total, 4)
+        out[mode] = {"total": total, **metrics}
+    return out
+
+
+def main(per_segment=40, output="docs/NCM_RETRIEVAL_EVALUATION_V2.json") -> dict:
+    session = SessionLocal()
+    samples = _samples(session, per_segment)
+    corpus = EnrichedNcmCorpus.build(session)
+    index = KbNeighborIndex.build(session)
+    lexicon = derive_abbreviation_lexicon(index.corpus(), min_support=3)
+    runs = {"OFFICIAL_ONLY": [], "KB_ONLY": [], "HYBRID": []}
+    for sample in samples:
+        # leave-one-out: the queried product never retrieves itself from the KB
+        runs["OFFICIAL_ONLY"].append(hierarchical_candidates(corpus, sample["description"], k=50))
+        runs["KB_ONLY"].append(kb_neighbor_candidates(index, sample["description"], k=50, exclude_product_id=sample["product_id"]))
+        runs["HYBRID"].append(hybrid_candidates(corpus, index, sample["description"], k=50, lexicon=lexicon, exclude_product_id=sample["product_id"]))
+    overall = _recall_mrr(runs, samples)
+    by_quality = {quality: _recall_mrr({mode: [row for row, sample in zip(list_rows, samples) if sample["quality"] == quality] for mode, list_rows in runs.items()},
+                                       [sample for sample in samples if sample["quality"] == quality]) for quality in ("GOLD", "SILVER", "NOISY")}
+    by_segment = {segment: _recall_mrr({mode: [row for row, sample in zip(list_rows, samples) if sample["segment"] == segment] for mode, list_rows in runs.items()},
+                                        [sample for sample in samples if sample["segment"] == segment]) for segment in sorted({sample["segment"] for sample in samples})}
+    report = {
+        "artifact": "NCM_RETRIEVAL_EVALUATION_V2",
+        "eval_total": len(samples),
+        "quality_counts": {quality: sum(1 for sample in samples if sample["quality"] == quality) for quality in ("GOLD", "SILVER", "NOISY")},
+        "lexicon_size": len(lexicon),
+        "lexicon_sample": dict(sorted(lexicon.items())[:40]),
+        "overall": overall,
+        "by_quality": by_quality,
+        "by_segment": by_segment,
+        "llm_liberation_gate": {"gold_recall_at_20_min": 0.80, "gold_recall_at_50_preferred": 0.90},
+        "cost_usd": 0.0,
+    }
+    Path(output).write_text(json.dumps(report, indent=2, ensure_ascii=False), encoding="utf-8")
+    session.close()
+    return report
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--per-segment", type=int, default=40)
+    parser.add_argument("--output", default="docs/NCM_RETRIEVAL_EVALUATION_V2.json")
+    args = parser.parse_args()
+    print(json.dumps(main(args.per_segment, args.output), indent=2))