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.

Example of Custom Policy

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):
    ...

Example of Custom NLU Component

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.

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:
        # Implementation details here
        ...

Registration

To make your graph component available to Rasa you may have to register your graph component with a recipe. Rasa uses recipes to translate the content of your model configuration to executable graphs.

from rasa.engine.graph import GraphComponent
from rasa.engine.recipes.default_recipe import DefaultV1Recipe

@DefaultV1Recipe.register(
    [DefaultV1Recipe.ComponentType.INTENT_CLASSIFIER], is_trainable=True
)
class MyComponent(GraphComponent):
    ...

Using Custom Components in your Model Configuration

You can use custom graph components like any other NLU component or policy within your model configuration. The only change is that you have to specify the full module name instead of the class name only.

recipe: default.v1
language: en
pipeline:
  - name: your.custom.NLUComponent
    setting_a: 0.01
    setting_b: string_value
policies:
  - name: your.custom.Policy

Implementation Hints

... ...