tonga

A lightweight extraction of the config-driven object construction system from ai2-tango. Provides FromParams, Registrable, Lazy, and Params with minimal dependencies (just rjsonnet), removing all Step/Workspace/integration machinery.

Installation

pip install tonga-config

The distribution name on PyPI is tonga-config, but the import name is tonga:

from tonga import FromParams, Registrable, Lazy, Params

For development:

pip install -e ".[dev]"

Usage

Registrable classes

Define a base class and register subclasses by name. Instantiate them from config dicts using from_params:

from tonga import Registrable, Params

class Encoder(Registrable):
    pass

@Encoder.register("lstm")
class LSTMEncoder(Encoder):
    def __init__(self, input_dim: int, hidden_dim: int, num_layers: int = 1):
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.num_layers = num_layers

encoder = Encoder.from_params(Params({
    "type": "lstm",
    "input_dim": 128,
    "hidden_dim": 256,
    "num_layers": 2,
}))

Loading configs from jsonnet files

from tonga import Params

params = Params.from_file("config.jsonnet")
model = Model.from_params(params)

Nested construction

Constructor arguments that are themselves FromParams/Registrable subclasses are constructed recursively:

from tonga import Registrable, Params

class Activation(Registrable):
    pass

@Activation.register("relu")
class ReLU(Activation):
    pass

@Activation.register("gelu")
class GELU(Activation):
    pass

class FeedForward(Registrable):
    pass

@FeedForward.register("linear")
class Linear(FeedForward):
    def __init__(self, input_dim: int, output_dim: int, activation: Activation):
        self.input_dim = input_dim
        self.output_dim = output_dim
        self.activation = activation

ff = FeedForward.from_params(Params({
    "type": "linear",
    "input_dim": 128,
    "output_dim": 64,
    "activation": {"type": "gelu"},
}))
assert isinstance(ff.activation, GELU)

Lazy construction

Use Lazy when one constructor argument depends on another:

from tonga import Registrable, Lazy, Params

class Optimizer(Registrable):
    pass

@Optimizer.register("sgd")
class SGD(Optimizer):
    def __init__(self, lr: float, params: list):
        self.lr = lr
        self.params = params

class Trainer(Registrable):
    pass

@Trainer.register("default")
class DefaultTrainer(Trainer):
    @classmethod
    def from_parts(cls, model_params: list, lazy_optimizer: Lazy[Optimizer]):
        trainer = cls()
        trainer.optimizer = lazy_optimizer.construct(params=model_params)
        return trainer

FromParams without registration

Classes can use FromParams directly without a registry when you don't need polymorphic dispatch:

from tonga import FromParams

class Config(FromParams):
    def __init__(self, learning_rate: float, batch_size: int = 32):
        self.learning_rate = learning_rate
        self.batch_size = batch_size

config = Config.from_params({"learning_rate": 0.001})
assert config.batch_size == 32

Acknowledgements

Tonga is a derivative work of ai2-tango by the Allen Institute for Artificial Intelligence. The FromParams, Registrable, Lazy, and Params classes were extracted and adapted from that project. All credit for the original design goes to the AI2 Tango authors and, before them, the AllenNLP team.

License

Tonga is licensed under the Apache License, Version 2.0. See the NOTICE file for upstream attribution.