ONNXRuntime

class ONNXRuntime(config)

Bases: Runtime

Step that process data using models in the ONNX format (https://onnx.ai/).

Input - the input fields need to be of type that is expected by the provided model.

Output - the output fields depend on the provided model and the data it returns.

Parameters

• type (str) – onnx

• model (str) – Model location

• input_mapping (dict) – dictionary of input fields mapping: name of the model input to the document input field

• output_mapping (dict) – dictionary of output fields mapping: name of the model output to the document output field

Example

{
    "step": "runtime",
    "type": "onnx",
    "model": "/models/test-model.onnx",
    "input_mapping": {
        "input": "query"
    },
    "output_mapping": {
        "output_label": "query_class",
        "output_probability": "query_probability"
    }
}

Methods Summary

pre_session_hook()	Pre-session hook that is executed before creating an inference session.
process_batch(batch)	Process a batch of documents.

Methods Documentation

pre_session_hook()

Pre-session hook that is executed before creating an inference session.

It can be overwritten by subclasses to perform custom operations e.g. downloading model.

Return type

None

process_batch(batch)

Process a batch of documents. If not defined will default to using self.process_doc for each document in the batch.

Parameters

batch (list(Document)) – List of documents

Returns

List of processed documents

Return type

list(Document)