Custom Graph Components

You are viewing documentation for our open source project which is maintained by the community. If you want to get started building assistants with Rasa please check out our latest documentation here.

Rasa provides a variety of NLU components and policies out of the box. You can customize them or create your own components from scratch by using custom graph components.

To use your custom graph component with Rasa it has to fulfill the following requirements:

• It has to implement the GraphComponent interface
• It has to be registered with the used model configuration
• It has to be used in the configuration file
• It has to use type annotations. Rasa makes use of the type annotations to validate your model configuration. Forward references are not allowed. If you're using Python 3.7 you can use from __future__ import annotations to get rid of forward references.

Graph Components

Rasa uses the passed in model configuration to build a directed acyclic graph. This graph describes the dependencies between the items in your model configuration and how data flows between them. This has two major benefits:

• Rasa can use the computational graph to optimize the execution of your model. Examples for this are efficient caching of training steps or executing independent steps in parallel.
• Rasa can represent different model architectures flexibly. As long as the graph remains acyclic Rasa can in theory pass any data to any graph component based on the model configuration without having to tie the underlying software architecture to the used model architecture.

When translating the model configuration to the computational graph policies and NLU components become nodes within this graph. While there is a distinction between policies and NLU components in your model configuration, the distinction is abstracted away when they are placed within the graph. At this point, policies and NLU components become abstract graph components. In practice this is represented by the GraphComponent interface: Both policies and NLU components have to inherit from this interface to become compatible and executable for Rasa's graph.

Getting Started

Before you get started, you have to decide whether you want to implement a custom NLU component or a policy. If you are implementing a custom policy, then we recommend extending the existing rasa.core.policies.policy.Policy class which already implements the GraphComponent interface.

from rasa.core.policies.policy import Policy

from rasa.engine.recipes.default_recipe import DefaultV1Recipe

# TODO: Correctly register your graph component
@DefaultV1Recipe.register(
    [DefaultV1Recipe.ComponentType.POLICY_WITHOUT_END_TO_END_SUPPORT], is_trainable=True
)
class MyPolicy(Policy):
    ...

If you want to implement a custom NLU component then start out with the following skeleton:

from typing import Dict, Text, Any, List

from rasa.engine.graph import GraphComponent, ExecutionContext
from rasa.engine.recipes.default_recipe import DefaultV1Recipe
from rasa.engine.storage.resource import Resource
from rasa.engine.storage.storage import ModelStorage
from rasa.shared.nlu.training_data.message import Message
from rasa.shared.nlu.training_data.training_data import TrainingData

# TODO: Correctly register your component with its type
@DefaultV1Recipe.register(
    [DefaultV1Recipe.ComponentType.INTENT_CLASSIFIER], is_trainable=True
)
class CustomNLUComponent(GraphComponent):
    @classmethod
    def create(cls,
        config: Dict[Text, Any],
        model_storage: ModelStorage,
        resource: Resource,
        execution_context: ExecutionContext,
    ) -> GraphComponent:
        # TODO: Implement this
        ...

def train(self, training_data: TrainingData) -> Resource:
        # TODO: Implement this if your component requires training
        ...

def process_training_data(self, training_data: TrainingData) -> TrainingData:
        # TODO: Implement this if your component augments the training data with
        # tokens or message features which are used by other components
        # during training.
        ...
        return training_data

def process(self, messages: List[Message]) -> List[Message]:
        # TODO: This is the method which Rasa Open Source will call during inference.
        ...
        return messages

Read the following sections to find out how to solve the TODOs in the example above and what other methods need to be implemented in your custom component.

Custom Tokenizers

If you create a custom tokenizer, you should extend the rasa.nlu.tokenizers.tokenizer.Tokenizer class. The train and process methods are already implemented so you only need to overwrite the tokenize method.

The GraphComponent Interface

To run your custom NLU component or policy with Rasa it must implement the GraphComponent interface.

from __future__ import annotations
from abc import ABC, abstractmethod
from typing import List, Type, Dict, Text, Any, Optional
from rasa.engine.graph import ExecutionContext
from rasa.engine.storage.resource import Resource
from rasa.engine.storage.storage import ModelStorage

class GraphComponent(ABC):
    """Interface for any component which will run in a graph."""

@classmethod
    def required_components(cls) -> List[Type]:
        """Components that should be included in the pipeline before this component."""
        return []

@classmethod
    @abstractmethod
    def create(
        cls,
        config: Dict[Text, Any],
        model_storage: ModelStorage,
        resource: Resource,
        execution_context: ExecutionContext,
    ) -> GraphComponent:
        """Creates a new `GraphComponent`."""
        ...

@classmethod
    def load(
        cls,
        config: Dict[Text, Any],
        model_storage: ModelStorage,
        resource: Resource,
        execution_context: ExecutionContext,
        **kwargs: Any,
    ) -> GraphComponent:
        """Creates a component using a persisted version of itself."""
        return cls.create(config, model_storage, resource, execution_context)

@staticmethod
    def get_default_config() -> Dict[Text, Any]:
        """Returns the component's default config."""
        return {}

@staticmethod
    def supported_languages() -> Optional[List[Text]]:
        """Determines which languages this component can work with."""
        return None

@staticmethod
    def not_supported_languages() -> Optional[List[Text]]:
        """Determines which languages this component cannot work with."""
        return None

@staticmethod
    def required_packages() -> List[Text]:
        """Any extra python dependencies required for this component to run."""
        return []

@classmethod
    def fingerprint_addon(cls, config: Dict[str, Any]) -> Optional[str]:
        """Adds additional data to the fingerprint calculation."""
        return None

Model Persistence

Some graph components require persisting data during training which should be available to the graph component at inference time. A typical use case is storing model weights. Rasa provides the model_storage and resource parameters to your graph component's create and load method for this purpose as shown in the snippet below:

from __future__ import annotations
from typing import Any, Dict, Text
from rasa.engine.graph import GraphComponent, ExecutionContext
from rasa.engine.storage.resource import Resource
from rasa.engine.storage.storage import ModelStorage

class MyComponent(GraphComponent):
    def __init__(self,
        model_storage: ModelStorage,
        resource: Resource,
        training_artifact: Optional[Dict],
    ) -> None:
        # Store both `model_storage` and `resource` as object attributes to be able
        # to utilize them at the end of the training
        self._model_storage = model_storage
        self._resource = resource

@classmethod
    def create(
        cls,
        config: Dict[Text, Any],
        model_storage: ModelStorage,
        resource: Resource,
        execution_context: ExecutionContext,
    ) -> MyComponent:
        return cls(model_storage, resource, training_artifact=None)

def train(self, training_data: TrainingData) -> Resource:
        # Train your graph component
        ...
        # Persist your graph component
        with self._model_storage.write_to(self._resource) as directory_path:
            with open(directory_path / "artifact.json", "w") as file:
                json.dump({"my": "training artifact"}, file)
        # Return resource to make sure the training artifacts can be cached.
        return self._resource

Writing to the Model Storage

The snippet below illustrates how to write your graph component's data to the model storage. To persist your graph component after training, the train method will need to access the values of model_storage and resource. Therefore, you should store the values of model_storage and resource at initialization time.

Reading from the Model Storage

Rasa will call the load method of your graph component to instantiate it for inference. You can use the context manager self._model_storage.read_from(resource) to get a path to the directory where your graph component's data was persisted. Using the provided path you can then load the persisted data and initialize your graph component with it. Note that the model_storage will throw a ValueError in case no persisted data was found for the given resource.