Files
duocthu/apps/ai-service/rag/run_eval.py
T

98 lines
3.2 KiB
Python

from __future__ import annotations
import argparse
import json
from pathlib import Path
from .artifacts import build_drug_catalog, load_aliases, load_documents, load_parents
from .evaluation import CaseOrigin, EvaluationCase, EvaluationOutcome, summarize
from .in_memory import InMemoryLexicalRetriever, InMemoryParentStore
from .models import EvidenceDecision, QueryIntent, SubjectScope
from .routing import CatalogDrugResolver, QueryRoutingService
from .service import EvidencePolicy, RetrievalService
def read_cases(path: Path) -> list[EvaluationCase]:
with path.open(encoding="utf-8") as handle:
return [
EvaluationCase(
case_id=raw["case_id"],
query=raw["query"],
expected_drug_id=raw.get("expected_drug_id"),
expected_id=raw.get("expected_id"),
origin=CaseOrigin(raw["origin"]),
subject_scope=SubjectScope(raw.get("subject_scope", "human")),
)
for line in handle
if line.strip()
for raw in [json.loads(line)]
]
def run(
cases_path: Path,
documents_path: Path,
parents_path: Path,
aliases_path: Path | None = None,
) -> dict:
documents = load_documents(documents_path)
retrieval = RetrievalService(
InMemoryLexicalRetriever(documents),
InMemoryParentStore(load_parents(parents_path)),
EvidencePolicy(),
)
service = QueryRoutingService(
retrieval,
CatalogDrugResolver(build_drug_catalog(documents, load_aliases(aliases_path))),
)
outcomes = []
details = []
for case in read_cases(cases_path):
result = service.retrieve(
case.query,
case.subject_scope,
QueryIntent.FACT_LOOKUP,
)
retrieved = (
tuple(item.evidence_id for item in result.evidence)
if result.decision != EvidenceDecision.ABSTAIN
else ()
)
outcome = EvaluationOutcome(
case=case,
retrieved_ids=retrieved,
resolved_drug_id=result.resolved_drug_id,
drug_resolution_status=result.drug_resolution_status,
)
outcomes.append(outcome)
details.append({
"case_id": case.case_id,
"passed": outcome.passed,
"decision": result.decision,
"reason": result.reason,
"expected_id": case.expected_id,
"retrieved_ids": retrieved,
"expected_drug_id": case.expected_drug_id,
"resolved_drug_id": result.resolved_drug_id,
"drug_resolution_status": result.drug_resolution_status,
})
return {**summarize(outcomes), "details": details}
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--cases", type=Path, required=True)
parser.add_argument("--documents", type=Path, required=True)
parser.add_argument("--parents", type=Path, required=True)
parser.add_argument("--aliases", type=Path)
args = parser.parse_args()
print(json.dumps(
run(args.cases, args.documents, args.parents, args.aliases),
ensure_ascii=False,
indent=2,
))
if __name__ == "__main__":
main()