predictor.py

"""Run a trained model on a document and get back decoded predictions.

Port of the one thing `radgraph.dygie.predictors.dygie.DyGIEPredictor` did that this project
actually used: forward pass -> `PredictedNER` / `PredictedRelation` objects, in document
(here: sentence, since there's always exactly one) coordinates.
"""
import torch

from .dataset import Example, build_example
from .document import Document
from .model import DyGIEModel
from .vocab import Vocabulary


@torch.no_grad()
def predict_document(model: DyGIEModel, vocab: Vocabulary, doc: Document, max_span_width: int) -> dict:
    model.eval()
    example: Example = build_example(doc, vocab, max_span_width)
    output = model(example)
    return {
        "doc_key": doc.doc_key,
        "example": example,
        "predicted_ner": output["ner"].get("predictions", []),
        "predicted_relations": output["relation"].get("predictions", []),
    }