116 lines
6.0 KiB
Python
116 lines
6.0 KiB
Python
"""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))
|