🧅 Unions

To allow for a union of types, you can use the protobuf oneof keyword.

// oneof_demo.proto
syntax = "proto3";

package oneof_demo;

// Transformer configuration
message Transformer {
    // Number of layers
    uint32 num_layers = 1;
    // Number of heads
    uint32 num_heads = 2;
    // Activation function
    string activation = 3;
}

// Conv block
message ConvBlock {
    // Number of output channels
    uint32 out_channels = 1;
    // Kernel size
    uint32 kernel_size = 2;
    // Activation function
    string activation = 3;
}

// Convolutional neural network configuration
message ConvNet {
    // Conv layer configuration
    repeated ConvBlock layers = 1;
}

// Model configuration
message Model {
    oneof backbone {
        Transformer transformer = 1;
        ConvNet conv_net = 2;
    }
}

The generated code will look like this:

# Autogenerated code. DO NOT EDIT.
import typing
import py_gen_ml as pgml


class Transformer(pgml.YamlBaseModel):
    """Transformer configuration"""

    num_layers: int
    """Number of layers"""

    num_heads: int
    """Number of heads"""

    activation: str
    """Activation function"""


class ConvBlock(pgml.YamlBaseModel):
    """Conv block"""

    out_channels: int
    """Number of output channels"""

    kernel_size: int
    """Kernel size"""

    activation: str
    """Activation function"""


class ConvNet(pgml.YamlBaseModel):
    """Convolutional neural network configuration"""

    layers: typing.List[ConvBlock]
    """Conv layer configuration"""


class Model(pgml.YamlBaseModel):
    """Model configuration"""

    backbone: typing.Union[Transformer, ConvNet]