| from __future__ import annotations |
|
|
| import argparse |
| from collections import defaultdict |
| import glob |
| import json |
| from pathlib import Path |
| import sys |
| from typing import Any |
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parents[1])) |
|
|
| import numpy as np |
|
|
| from patientphex.ontology import HPOOntology |
| from patientphex.sapbert import _unique_alias_records, encode_cls, rank_concepts |
|
|
|
|
| def load_jsonl(path: Path) -> list[dict[str, Any]]: |
| return [json.loads(line) for line in path.open(encoding="utf-8") if line.strip()] |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser(description="Ensemble verifier scores and normalize selected A-set spans.") |
| parser.add_argument("--data", type=Path, required=True) |
| parser.add_argument("--ontology", type=Path, required=True) |
| parser.add_argument("--verifier-glob", required=True) |
| parser.add_argument("--sapbert-model", type=Path, required=True) |
| parser.add_argument("--output", type=Path, required=True) |
| parser.add_argument("--threshold", type=float, default=0.66) |
| parser.add_argument( |
| "--target-fraction", type=float, |
| help="If supplied, replace the raw threshold by the ensemble-score quantile that keeps this fraction.", |
| ) |
| parser.add_argument("--lexicon-threshold", type=float, default=0.0) |
| parser.add_argument("--phenotagger-threshold", type=float, default=0.90) |
| parser.add_argument("--device", default="cuda") |
| args = parser.parse_args() |
|
|
| paths = sorted(glob.glob(args.verifier_glob)) |
| if not paths: |
| raise FileNotFoundError(f"no verifier outputs matched {args.verifier_glob!r}") |
| grouped: dict[tuple[str, int, int], list[dict[str, Any]]] = defaultdict(list) |
| for path in paths: |
| result = json.loads(Path(path).read_text(encoding="utf-8")) |
| for record in result["predictions"]: |
| grouped[(str(record["pmc_id"]), int(record["start"]), int(record["end"]))].append(record) |
| expected_models = len(paths) |
| incomplete = [key for key, records in grouped.items() if len(records) != expected_models] |
| if incomplete: |
| raise RuntimeError(f"{len(incomplete)} candidates are missing ensemble scores") |
| ensemble_records: list[dict[str, Any]] = [] |
| for records in grouped.values(): |
| record = dict(records[0]) |
| record["ensemble_score"] = float(np.mean([float(item["score"]) for item in records])) |
| ensemble_records.append(record) |
| threshold = args.threshold |
| if args.target_fraction is not None: |
| if not 0.0 < args.target_fraction <= 1.0: |
| parser.error("--target-fraction must be in (0, 1]") |
| threshold = float(np.quantile( |
| [record["ensemble_score"] for record in ensemble_records], 1.0 - args.target_fraction |
| )) |
| selected = [record for record in ensemble_records if record["ensemble_score"] >= threshold] |
|
|
| ontology = HPOOntology.from_obo(args.ontology) |
|
|
| def valid_identifier(value: Any) -> str | None: |
| if value is None: |
| return None |
| identifier = ontology.canonical_id(str(value)) |
| return identifier if ontology.classify_id(identifier) == "valid_branch_hp" else None |
|
|
| unresolved = sorted({ |
| str(record["text"]).lower() for record in selected |
| if not ( |
| valid_identifier(record.get("lexicon_identifier")) is not None |
| and float(record.get("lexicon_quality", 0.0)) >= args.lexicon_threshold |
| ) and not ( |
| valid_identifier(record.get("phenotagger_identifier")) is not None |
| and float(record.get("phenotagger_score", 0.0)) >= args.phenotagger_threshold |
| ) |
| }) |
| ranking_by_text: dict[str, list[str]] = {} |
| if unresolved: |
| from transformers import AutoModel, AutoTokenizer |
| tokenizer = AutoTokenizer.from_pretrained(args.sapbert_model, local_files_only=True) |
| model = AutoModel.from_pretrained(args.sapbert_model, local_files_only=True).to(args.device) |
| alias_ids, alias_texts = _unique_alias_records(ontology) |
| alias_embeddings = encode_cls(alias_texts, tokenizer, model, device=args.device, batch_size=256) |
| query_embeddings = encode_cls(unresolved, tokenizer, model, device=args.device, batch_size=256) |
| rankings = rank_concepts(query_embeddings, alias_embeddings, alias_ids, device=args.device, max_concepts=5) |
| ranking_by_text = dict(zip(unresolved, rankings, strict=True)) |
|
|
| entities_by_doc: dict[str, dict[tuple[int, int, str], dict[str, Any]]] = defaultdict(dict) |
| source_counts: dict[str, int] = defaultdict(int) |
| for record in selected: |
| lexicon_identifier = valid_identifier(record.get("lexicon_identifier")) |
| phenotagger_identifier = valid_identifier(record.get("phenotagger_identifier")) |
| if ( |
| lexicon_identifier is not None |
| and float(record.get("lexicon_quality", 0.0)) >= args.lexicon_threshold |
| ): |
| identifier = lexicon_identifier |
| source = "lexicon" |
| elif ( |
| phenotagger_identifier is not None |
| and float(record.get("phenotagger_score", 0.0)) >= args.phenotagger_threshold |
| ): |
| identifier = phenotagger_identifier |
| source = "phenotagger" |
| else: |
| identifier = ranking_by_text[str(record["text"]).lower()][0] |
| source = "sapbert" |
| source_counts[source] += 1 |
| start, end = int(record["start"]), int(record["end"]) |
| entities_by_doc[str(record["pmc_id"])][(start, end, identifier)] = { |
| "identifier": identifier, "type": "Phenotype", "offset": start, |
| "length": end - start, "text": record["text"], "note": None, |
| } |
|
|
| documents = load_jsonl(args.data) |
| args.output.parent.mkdir(parents=True, exist_ok=True) |
| with args.output.open("w", encoding="utf-8", newline="\n") as handle: |
| for document in documents: |
| entities = sorted( |
| entities_by_doc[str(document["pmc_id"])].values(), |
| key=lambda item: (int(item["offset"]), int(item["length"]), str(item["identifier"])), |
| ) |
| output = { |
| "pmc_id": document["pmc_id"], "pmid": document.get("pmid"), |
| "entities": entities, "association": [], |
| } |
| handle.write(json.dumps(output, ensure_ascii=False, separators=(",", ":")) + "\n") |
| print(json.dumps({ |
| "models": expected_models, "threshold": threshold, "target_fraction": args.target_fraction, |
| "selected_candidates": len(selected), "entities": sum(len(x) for x in entities_by_doc.values()), |
| "normalization_sources": dict(source_counts), "output": str(args.output), |
| })) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
