File size: 6,703 Bytes
e7543dc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 | 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()
|