Use ML models in your worker

Overview

Arkindex workers can use AI models to process documents. In most cases, we want to store them on Arkindex for the following reasons:

  • Model files, which can be very large, should not be stored directly in the worker Docker image,

  • If multiple Ponos agents are deployed on the same machine, Arkindex can cache the files (this means less downloads),

  • Workers can be compatible with multiple models, we want to let the user pick what model they want to use and Arkindex has an interface for that.

Publication on Arkindex

There are two ways to publish models on Arkindex. Both require the usage of Arkindex’s CLI tool.

Make sure that the file architecture is compatible with what is described in the worker’s documentation.

Avoid unnecessary folders in the model files to make sure that the model loading in the worker is smooth.

How to load them in your worker

To tell Arkindex that the worker needs access to Arkindex models, set the model_usage parameter in the YAML configuration to either:

  • supported: the worker can run both with or without an Arkindex model,

  • required: the worker cannot run without an Arkindex model.

The model loading should be done once per worker execution, as this can take some time. We recommend implementing that in the configure method.

class MyWorker(ElementsWorker):
    def configure(self):
        super().configure()

        # Model loading code
        self.load_model()

    def load_model(self):
        # Look for the folder where the model files are stored
        model_directory = self.find_extras_directory()

        # Instantiate the model
        self.model = ...

The Ponos agent will take care of downloading the model. The path towards the download files is available through the find_extras_directory method.

When running the worker locally, the result of find_extras_directory is the path given through the CLI argument --extras-dir.
This path links to a folder, so if needed, you still need to iterate over its content to find the right file to load.