model to export

optional model task or category

optional output model path, defaults to return byte array if None

onnx opset, defaults to 12

path to model output or model as bytes depending on output parameter

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

def call(self, model, task=”default”, output=None, opset=12):
    “””
    Exports a machine learning model to ONNX using ONNXMLTools.
    Args:
        model: model to export
        task: optional model task or category
        output: optional output model path, defaults to return byte array if None
        opset: onnx opset, defaults to 12
    Returns:
        path to model output or model as bytes depending on output parameter
    “””
    # Convert scikit-learn model to ONNX
    model = convert_sklearn(model, task, initial_types=[(“input_ids”, StringTensorType([None, None]))], target_opset=opset)
    # Prune model graph down to only output probabilities
    model = select_model_inputs_outputs(model, outputs=”probabilities”)
    # pylint: disable=E1101
    # Rename output to logits for consistency with other models
    model.graph.output[0].name = “logits”
    # Find probabilities output node and rename to logits
    for node in model.graph.node:
        if node.output[0] == “probabilities”:
            node.output[0] = “logits”
    # Save model to specified output path or return bytes
    model = save_onnx_model(model, output)
    return output if output else model