inference_dcpg.py

from __future__ import annotations

import json
import sys

from dcpg_encoder import DCPGEncoder, encode_patient

_encoder = DCPGEncoder()


def predict(graph_summary: dict, source: str = "dcpg") -> dict:
    return encode_patient(graph_summary, encoder=_encoder, source=source)


def predict_batch(summaries: list, source: str = "dcpg") -> list:
    return [predict(s, source=source) for s in summaries]


if __name__ == "__main__":
    if len(sys.argv) > 1:
        with open(sys.argv[1]) as f:
            data = json.load(f)
        result = predict(data)
    else:
        result = predict({
            "nodes": [
                {"node_id": "p1::text::NAME_DATE_MRN_FACILITY", "modality": "text",
                 "phi_type": "NAME_DATE_MRN_FACILITY", "risk_entropy": 0.8,
                 "context_confidence": 0.9, "pseudonym_version": 2},
                {"node_id": "p1::audio_proxy::VOICE", "modality": "audio_proxy",
                 "phi_type": "VOICE", "risk_entropy": 0.55,
                 "context_confidence": 0.6, "pseudonym_version": 1},
            ],
            "edges": [
                {"source": "p1::text::NAME_DATE_MRN_FACILITY",
                 "target": "p1::audio_proxy::VOICE",
                 "type": "cross_modal", "weight": 0.63},
            ],
        })
    print(json.dumps(result, indent=2))