<!-- Source: https://docs.biosimulant.com/how-to/add-onnx-model -->

# How to Add an ONNX Model

Use ONNX for finite inference or repeated inference inside the same
communication-step world as mechanistic modules.

## Step 1: export the model

PyTorch:

```python

model = MyClassifier()
model.eval()
dummy_input = torch.randn(1, 10)
torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    input_names=["input"],
    output_names=["probabilities"],
)
```

## Step 2: wrap it as a BioModule

```python

class ONNXClassifier(biosim.BioModule):
    execution_policy = biosim.ExecutionPolicy.ONCE_BEFORE_RUN

    def __init__(self, model_path: str = "data/assets/model.onnx"):
        import onnxruntime as ort

        self.session = ort.InferenceSession(model_path)
        self.input_name = self.session.get_inputs()[0].name
        self.output_name = self.session.get_outputs()[0].name

    def inputs(self):
        return {
            "state_vector": biosim.SignalSpec.array(
                dtype="float32",
                shape=(10,),
            )
        }

    def outputs(self):
        return {
            "classification": biosim.SignalSpec.scalar(dtype="int64"),
            "probabilities": biosim.SignalSpec.array(dtype="float32", shape=(3,)),
        }

    def execute(self, inputs, *, context):
        input_data = np.asarray(inputs["state_vector"].value, dtype=np.float32).reshape(1, -1)
        result = self.session.run([self.output_name], {self.input_name: input_data})
        prediction = np.asarray(result[0][0], dtype=np.float32)
        class_idx = int(np.argmax(prediction))
        return {
            "classification": class_idx,
            "probabilities": prediction.tolist(),
        }
```

## Step 3: declare the manifest

```yaml
schema_version: "2.0"
title: "Neuron State Classifier"
description: "ONNX classifier for neuron state vectors"
standard: onnx

biosim:
  entrypoint: "src.onnx_classifier:ONNXClassifier"
  init_kwargs:
    model_path: "data/assets/model.onnx"
  communication_step: 0.001

io:
  inputs:
    - name: state_vector
      signal_type: array
      dtype: float32
      shape: [10]
  outputs:
    - name: classification
      signal_type: scalar
      dtype: int64
    - name: probabilities
      signal_type: array
      dtype: float32
      shape: [3]

onnx:
  task: classification
  model_file: data/assets/model.onnx
  inputs:
    - name: input
      dtype: float32
      shape: [-1, 10]
  outputs:
    - name: probabilities
      dtype: float32
      shape: [-1, 3]

runtime:
  python_version: "3.12"
  dependencies:
    packages:
      - onnxruntime>=1.16
      - numpy>=1.24
```

## Step 4: compose it in a lab

```yaml
models:
  - package: biosimulant/neuro-mechanistic
    version: 1.0.0
    alias: neuron

  - path: ../models/state-classifier
    alias: classifier

wiring:
  - from: neuron.state_vector
    to:
      - classifier.state_vector

runtime:
  duration: 0.05
  communication_step: 0.001
```

**Info:**

  Keep `ONCE_BEFORE_RUN` for finite classification or a pure inference pipeline.
  Use `ONCE_AFTER_RUN` for final analysis, or `EACH_WINDOW` when the ONNX model
  must consume evolving simulation state. Lab and model manifests do not add an
  execution-mode field.

## Next steps

- [Wrap an External Simulator](/how-to/wrap-external-simulator)
- [model.yaml Schema](/references/model-manifest)
- [BioModule API](/references/biomodule-api)
