Checkpoint frontend UI/UX overhaul and ingestion embed benchmark work

ingestion/ingestion/embed/benchmark_local.py

@@ -0,0 +1,247 @@
+"""Offline embedding benchmark aligned with deterministic drug resolution.
+
+The production RAG contract never lets vector similarity choose a drug. This
+benchmark therefore resolves the expected drug through the verified alias
+catalog, filters candidates to that drug, and measures whether an embedding
+retrieves the expected section. Policy/abstention and multi-drug cases are not
+silently converted into retrieval cases.
+
+Only BGE-M3 is accepted here. Cloud providers deliberately have no CLI switch;
+using one requires a separate, explicitly approved benchmark path.
+"""
+from __future__ import annotations
+
+import argparse
+import csv
+import json
+import math
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import Iterable, Mapping, Sequence
+
+from ingestion.entities.catalog import normalize_name
+
+from .local_bge_m3 import BgeM3Local
+from .ports import EmbeddingProvider
+
+
+ATTRIBUTE_TO_SECTION = {
+    "chong_chi_dinh": "chong_chi_dinh",
+    "lieu_dung": "lieu_luong_va_cach_dung",
+    "mang_thai": "thoi_ky_mang_thai",
+    "qua_lieu": "qua_lieu_va_xu_tri",
+    "tac_dung_phu": "tac_dung_khong_mong_muon",
+    "than_trong": "than_trong",
+    "tuong_tac": "tuong_tac_thuoc",
+}
+
+
+@dataclass(frozen=True)
+class BenchmarkCase:
+    case_id: str
+    query: str
+    drug_id: str
+    expected_section: str
+
+
+@dataclass(frozen=True)
+class CaseResult:
+    case_id: str
+    drug_id: str
+    expected_section: str
+    first_relevant_rank: int | None
+    top_chunk_id: str
+    top_section: str
+
+
+@dataclass(frozen=True)
+class BenchmarkReport:
+    model_id: str
+    dimensions: int
+    case_count: int
+    candidate_chunk_count: int
+    document_latency_ms: float
+    query_latency_ms: float
+    document_requests: int
+    query_requests: int
+    hit_at_1: float
+    hit_at_3: float
+    hit_at_5: float
+    mrr: float
+    cases: list[CaseResult]
+
+
+def _read_jsonl(path: Path) -> list[dict]:
+    with path.open(encoding="utf-8") as handle:
+        return [json.loads(line) for line in handle if line.strip()]
+
+
+def _alias_lookup(path: Path) -> dict[str, set[str]]:
+    payload = json.loads(path.read_text(encoding="utf-8"))
+    lookup: dict[str, set[str]] = {}
+    for entity in payload["entities"]:
+        for alias in entity["aliases"]:
+            key = normalize_name(alias)
+            lookup.setdefault(key, set()).add(entity["drug_id"])
+    return lookup
+
+
+def load_cases(golden_path: Path, entities_path: Path) -> list[BenchmarkCase]:
+    aliases = _alias_lookup(entities_path)
+    cases: list[BenchmarkCase] = []
+    with golden_path.open(encoding="utf-8-sig", newline="") as handle:
+        for row in csv.DictReader(handle):
+            attribute = row["thuoc_tinh_ky_vong"].strip()
+            drug_name = row["thuoc_ky_vong"].strip()
+            if not attribute or not drug_name or ";" in drug_name:
+                continue
+            section = ATTRIBUTE_TO_SECTION.get(attribute)
+            if section is None:
+                raise ValueError(
+                    f"golden case {row['id']} has unmapped attribute {attribute!r}"
+                )
+            drug_ids = aliases.get(normalize_name(drug_name), set())
+            if not drug_ids:
+                raise ValueError(
+                    f"golden case {row['id']} drug {drug_name!r} is not in the "
+                    "verified entity catalog"
+                )
+            if len(drug_ids) != 1:
+                raise ValueError(
+                    f"golden case {row['id']} drug {drug_name!r} is ambiguous: "
+                    f"{sorted(drug_ids)}"
+                )
+            drug_id = next(iter(drug_ids))
+            cases.append(
+                BenchmarkCase(
+                    case_id=row["id"],
+                    query=row["cau_hoi"].strip(),
+                    drug_id=drug_id,
+                    expected_section=section,
+                )
+            )
+    if not cases:
+        raise ValueError("golden dataset has no single-drug retrieval cases")
+    return cases
+
+
+def candidate_chunks(chunks_path: Path, cases: Sequence[BenchmarkCase]) -> list[dict]:
+    drug_ids = {case.drug_id for case in cases}
+    chunks = [
+        chunk
+        for chunk in _read_jsonl(chunks_path)
+        if chunk["drug_id"] in drug_ids and chunk["chunk_kind"] == "prose"
+    ]
+    available = {(chunk["drug_id"], chunk["section_key"]) for chunk in chunks}
+    missing = [
+        case.case_id
+        for case in cases
+        if (case.drug_id, case.expected_section) not in available
+    ]
+    if missing:
+        raise ValueError(f"golden cases have no canonical target chunks: {missing}")
+    return chunks
+
+
+def _cosine(left: Sequence[float], right: Sequence[float]) -> float:
+    numerator = sum(a * b for a, b in zip(left, right, strict=True))
+    left_norm = math.sqrt(sum(value * value for value in left))
+    right_norm = math.sqrt(sum(value * value for value in right))
+    if left_norm == 0 or right_norm == 0:
+        raise ValueError("embedding benchmark received a zero vector")
+    return numerator / (left_norm * right_norm)
+
+
+def run_benchmark(
+    provider: EmbeddingProvider,
+    cases: Sequence[BenchmarkCase],
+    chunks: Sequence[Mapping[str, object]],
+) -> BenchmarkReport:
+    document_batch = provider.embed_documents([str(chunk["text"]) for chunk in chunks])
+    query_batch = provider.embed_queries([case.query for case in cases])
+    case_results: list[CaseResult] = []
+
+    for case, query_vector in zip(cases, query_batch.vectors, strict=True):
+        ranked = sorted(
+            (
+                (_cosine(query_vector.values, vector.values), chunk)
+                for chunk, vector in zip(chunks, document_batch.vectors, strict=True)
+                if chunk["drug_id"] == case.drug_id
+            ),
+            key=lambda item: (-item[0], str(item[1]["chunk_id"])),
+        )
+        if not ranked:
+            raise ValueError(f"case {case.case_id} has no candidate chunks")
+        relevant_rank = next(
+            (
+                rank
+                for rank, (_, chunk) in enumerate(ranked, start=1)
+                if chunk["section_key"] == case.expected_section
+            ),
+            None,
+        )
+        top = ranked[0][1]
+        case_results.append(
+            CaseResult(
+                case_id=case.case_id,
+                drug_id=case.drug_id,
+                expected_section=case.expected_section,
+                first_relevant_rank=relevant_rank,
+                top_chunk_id=str(top["chunk_id"]),
+                top_section=str(top["section_key"]),
+            )
+        )
+
+    count = len(case_results)
+    ranks = [result.first_relevant_rank for result in case_results]
+
+    def hit_at(limit: int) -> float:
+        return sum(rank is not None and rank <= limit for rank in ranks) / count
+
+    return BenchmarkReport(
+        model_id=provider.model_id,
+        dimensions=provider.dimensions,
+        case_count=count,
+        candidate_chunk_count=len(chunks),
+        document_latency_ms=document_batch.latency_ms,
+        query_latency_ms=query_batch.latency_ms,
+        document_requests=document_batch.request_count,
+        query_requests=query_batch.request_count,
+        hit_at_1=hit_at(1),
+        hit_at_3=hit_at(3),
+        hit_at_5=hit_at(5),
+        mrr=sum(0.0 if rank is None else 1.0 / rank for rank in ranks) / count,
+        cases=case_results,
+    )
+
+
+def write_report(report: BenchmarkReport, path: Path) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(
+        json.dumps(asdict(report), ensure_ascii=False, indent=2) + "\n",
+        encoding="utf-8",
+    )
+
+
+def main(argv: Iterable[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(description="Local BGE-M3 subset benchmark")
+    parser.add_argument("--golden", type=Path, required=True)
+    parser.add_argument("--chunks", type=Path, required=True)
+    parser.add_argument("--entities", type=Path, required=True)
+    parser.add_argument("--out", type=Path, required=True)
+    parser.add_argument("--batch-size", type=int, default=8)
+    parser.add_argument("--device", default="cpu")
+    args = parser.parse_args(argv)
+
+    cases = load_cases(args.golden, args.entities)
+    chunks = candidate_chunks(args.chunks, cases)
+    report = run_benchmark(
+        BgeM3Local(batch_size=args.batch_size, device=args.device), cases, chunks
+    )
+    write_report(report, args.out)
+    print(json.dumps(asdict(report), ensure_ascii=False, indent=2))
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

ingestion/tests/test_embed_benchmark_local.py

@@ -0,0 +1,157 @@
+import csv
+import json
+from pathlib import Path
+
+import pytest
+
+from ingestion.embed.benchmark_local import (
+    BenchmarkCase,
+    candidate_chunks,
+    load_cases,
+    run_benchmark,
+)
+from ingestion.embed.local_bge_m3 import BgeM3Local
+
+
+def _write_csv(path: Path, rows: list[dict]) -> None:
+    with path.open("w", encoding="utf-8", newline="") as handle:
+        writer = csv.DictWriter(handle, fieldnames=list(rows[0]))
+        writer.writeheader()
+        writer.writerows(rows)
+
+
+def test_load_cases_uses_verified_aliases_and_skips_policy_and_multidrug(tmp_path):
+    entities = tmp_path / "entities.json"
+    entities.write_text(
+        json.dumps(
+            {
+                "entities": [
+                    {"drug_id": "paracetamol", "aliases": ["Paracetamol"]}
+                ]
+            }
+        ),
+        encoding="utf-8",
+    )
+    golden = tmp_path / "golden.csv"
+    _write_csv(
+        golden,
+        [
+            {
+                "id": "1",
+                "cau_hoi": "Liều Paracetamol?",
+                "thuoc_ky_vong": "Paracetamol",
+                "thuoc_tinh_ky_vong": "lieu_dung",
+            },
+            {
+                "id": "2",
+                "cau_hoi": "Uống gì?",
+                "thuoc_ky_vong": "",
+                "thuoc_tinh_ky_vong": "",
+            },
+            {
+                "id": "3",
+                "cau_hoi": "So sánh",
+                "thuoc_ky_vong": "Paracetamol; Ibuprofen",
+                "thuoc_tinh_ky_vong": "lieu_dung",
+            },
+        ],
+    )
+
+    assert load_cases(golden, entities) == [
+        BenchmarkCase(
+            case_id="1",
+            query="Liều Paracetamol?",
+            drug_id="paracetamol",
+            expected_section="lieu_luong_va_cach_dung",
+        )
+    ]
+
+
+def test_candidate_chunks_fails_when_ground_truth_is_missing(tmp_path):
+    path = tmp_path / "chunks.jsonl"
+    path.write_text(
+        json.dumps(
+            {
+                "drug_id": "drug",
+                "section_key": "chi_dinh",
+                "chunk_kind": "prose",
+                "text": "text",
+            }
+        )
+        + "\n",
+        encoding="utf-8",
+    )
+    cases = [BenchmarkCase("1", "dose", "drug", "lieu_luong_va_cach_dung")]
+
+    with pytest.raises(ValueError, match="no canonical target"):
+        candidate_chunks(path, cases)
+
+
+def test_load_cases_rejects_an_alias_ambiguous_in_the_verified_catalog(tmp_path):
+    entities = tmp_path / "entities.json"
+    entities.write_text(
+        json.dumps(
+            {
+                "entities": [
+                    {"drug_id": "drug_a", "aliases": ["Shared"]},
+                    {"drug_id": "drug_b", "aliases": ["Shared"]},
+                ]
+            }
+        ),
+        encoding="utf-8",
+    )
+    golden = tmp_path / "golden.csv"
+    _write_csv(
+        golden,
+        [
+            {
+                "id": "1",
+                "cau_hoi": "Liều Shared?",
+                "thuoc_ky_vong": "Shared",
+                "thuoc_tinh_ky_vong": "lieu_dung",
+            }
+        ],
+    )
+
+    with pytest.raises(ValueError, match="is ambiguous"):
+        load_cases(golden, entities)
+
+
+def test_benchmark_filters_by_drug_and_reports_first_relevant_rank():
+    vectors = {
+        "query": [1.0, 0.0],
+        "wrong section": [0.9, 0.1],
+        "correct section": [0.8, 0.2],
+        "other drug": [1.0, 0.0],
+    }
+    provider = BgeM3Local(encoder=lambda texts: [vectors[text] for text in texts])
+    provider._check_dimensions = lambda _values: None
+    cases = [BenchmarkCase("1", "query", "drug", "target")]
+    chunks = [
+        {
+            "chunk_id": "wrong",
+            "drug_id": "drug",
+            "section_key": "other",
+            "text": "wrong section",
+        },
+        {
+            "chunk_id": "right",
+            "drug_id": "drug",
+            "section_key": "target",
+            "text": "correct section",
+        },
+        {
+            "chunk_id": "leak",
+            "drug_id": "other",
+            "section_key": "target",
+            "text": "other drug",
+        },
+    ]
+
+    report = run_benchmark(provider, cases, chunks)
+
+    assert report.case_count == 1
+    assert report.cases[0].first_relevant_rank == 2
+    assert report.hit_at_1 == 0.0
+    assert report.hit_at_3 == 1.0
+    assert report.mrr == 0.5