How to Set Up a Rasa Development Environment in Sagemaker Studiolab with VS Code

As an experienced developer and early adopter of conversational AI technology, I‘ve witnessed firsthand the explosive growth of chatbots and virtual assistants in recent years. Across industries from healthcare to retail to financial services, organizations are leveraging natural language interfaces to automate customer support, streamline operations, and gain valuable insights from user interactions.

While platforms like Google‘s Dialogflow and Amazon Lex provide turn-key solutions for basic chatbots, many enterprises are turning to open-source frameworks for maximum flexibility and customizability. And one of the most powerful and popular open-source chatbot frameworks today is Rasa.

What is Rasa?

Rasa is an open-source conversational AI framework for building contextual assistants. Developed by the startup Rasa Technologies and backed by investors like Accel and Basis Set Ventures, Rasa provides a suite of tools for:

• Natural Language Understanding (NLU): Rasa NLU is a library for extracting structured data from user messages using machine learning. It identifies the intent behind a message and extracts relevant entities.

• Dialogue Management: Rasa Core is a stateful chatbot framework that uses ML-trained policies to decide what action a bot should take in a conversation based on context. It also includes tools for testing dialogue flows and generating training stories.

• Action Fulfillment: Rasa SDK allows integrating custom actions in Python that the bot can execute to fetch data from APIs, query databases, or perform side effects in the world.

• Deployment: Rasa X and Rasa server provide an end-to-end pipeline for deploying and continuously improving a trained assistant.

According to the 2021 Conversational AI Survey by Rasa, 66% of developers building AI assistants are using open source tools, and Rasa is the most popular open source framework with 51% adoption. High-profile companies using Rasa in production include Adobe, Deutsche Telekom, BMW, Airbus, and Lemonade.

Compared to a cloud-based Natural Language Understanding (NLU) service like Google‘s Dialogflow or Microsoft LUIS, Rasa provides several advantages for enterprises:

• Full control over data and models: Rasa runs entirely on your own infrastructure, so you keep full ownership and control of your training data, models, and deployment. This is critical for enterprises dealing with sensitive user data.

• Flexibility: With access to the full source code, Rasa can be customized, extended, and integrated to fit any use case. You have full control over your conversational AI stack.

• Vendor independence: As an open-source framework, Rasa mitigates the risk of vendor lock-in. You can run Rasa anywhere and port your models between deployment targets.

Of course, this power and flexibility comes with a tradeoff in ease-of-use. Rasa has a significant learning curve, and building a complete production assistant requires substantial development work. This is where having the right development environment becomes crucial.

Why Sagemaker Studio Lab for Rasa Development?

While you can certainly develop Rasa locally on your laptop, there are several reasons why using a cloud IDE like AWS Sagemaker Studio Lab can provide a better development experience:

1. Compute Power: Training NLU models and dialogue policies can be computationally intensive, especially with large datasets. Even a beefy laptop can struggle with memory and CPU constraints. Sagemaker Studio Lab provides on-demand access to powerful CPU and GPU instances to speed up model training.

2. Collaboration: A cloud-based development environment makes it easy to collaborate with teammates. You can share notebooks, code, and models without worrying about environment discrepancies.

3. Reproducibility: With Sagemaker Studio Lab, you can define your exact development environment, including dependencies and configurations, as code. This ensures that your setup is reproducible and avoids "works on my machine" issues.

4. Cost: Best of all, Sagemaker Studio Lab is completely free to use. You get access to compute instances comparable to a ml.t3.medium (2 vCPU, 4 GiB) for 12 hours at a time, with up to 15 GB of persistent storage.

The only catch is that Studio Lab does not come with an integrated development environment (IDE) out of the box. But with the help of code-server, an open-source project that runs VS Code on a remote server, we can create a full-featured IDE right in Sagemaker Studio Lab.

Setting Up The Development Environment

Step 1: Create a Conda Environment

First, let‘s create a fresh Conda environment to isolate our Rasa dependencies. Open a terminal in Studio Lab and run:

conda create --name rasa python==3.8
conda activate rasa

Now let‘s install the Rasa libraries:

pip install rasa

We‘ll also install the Rasa SDK for running custom actions:

pip install rasa-sdk

Finally, let‘s install code-server to set up our VS Code environment:

conda install -c conda-forge code-server

To make it easy to share and reproduce our environment, let‘s export our Conda configuration:

conda env export > environment.yml

Step 2: Launch a VS Code Server

With our environment set up, we‘re ready to launch a VS Code server. In a new Studio Lab terminal, run:

code-server --auth none --port 8080

This starts the code-server without authentication and binds it to port 8080.

To access the VS Code interface, we need to construct a URL like:

https://<your-studiolab-id>.studio.us-east-1.sagemaker.aws/studiolab/default/jupyter/proxy/8080/

Paste this URL into a new browser tab, and voilà, you have a full-fledged VS Code IDE running in Sagemaker Studio Lab!

Step 3: Initialize a Rasa Project

With our development environment ready, let‘s create a new Rasa project. In the VS Code terminal, run:

mkdir rasa-assistant
cd rasa-assistant
rasa init

This initializes a starter Rasa project with some example training data and configurations.

Let‘s take a look at the key files:

• config.yml: Defines the NLU and dialogue model configuration and policies.
• domain.yml: Specifies the assistant‘s intents, entities, slots, responses, forms, and actions.
• data/nlu.yml: Contains training examples for the NLU model.
• data/stories.yml: Defines the dialogue flows and training stories for the dialogue model.
• actions/actions.py: Implements custom action code.
• endpoints.yml: Configures the assistant‘s action server, tracker store, and other deployment endpoints.

Step 4: Write Dialogue and Training Data

Now comes the fun part – designing our assistant‘s conversational abilities! Let‘s sketch out a simple banking assistant that can check balances and transfer money.

First, let‘s define our intents and entities in domain.yml:

intents:
  - greet
  - check_balance
  - transfer_money
  - goodbye

entities:
  - account_type
  - amount

slots:
  account_type:
    type: text
  amount: 
    type: float  

responses:
  utter_greet:
    - text: Hi there! How can I help you today?
  utter_balance:
    - text: Your {account_type} account balance is {balance}.
  utter_transfer_complete:  
    - text: OK, I‘ve transferred {amount} from your {source_account} to your {dest_account}.
  utter_goodbye:
    - text: Goodbye! Let me know if you need anything else.

actions:
  - action_check_balance
  - action_transfer_money  

Here we‘ve defined our assistant‘s capabilities:

• It can recognize intents for greeting, checking balances, transferring money, and saying goodbye.
• It extracts entities for account types and amounts.
• It tracks the account type and amount in conversation slots.
• It has responses to handle the various user intents.
• It has two custom actions for fetching balance information and executing transfers.

Next, let‘s provide some training examples for the NLU model in data/nlu.yml:

nlu:
- intent: greet
  examples: |
    - hi
    - hello
    - hey
    - good morning

- intent: check_balance
  examples: |
    - what‘s my [checking](account_type) balance?
    - how much money do I have in my [savings](account_type) account  
    - check my [credit card](account_type) balance

- intent: transfer_money
  examples: |
    - transfer [$500](amount) from checking to savings
    - I want to move [100](amount) dollars to checking 
    - send [1000](amount) to my credit card from savings  

And let‘s define a simple story flow in data/stories.yml:

stories:
- story: check balance
  steps:
  - intent: greet
  - action: utter_greet
  - intent: check_balance
    entities:
    - account_type: checking
  - action: action_check_balance
  - action: utter_goodbye

- story: transfer money  
  steps:  
  - intent: transfer_money
    entities:
    - amount: 100
  - action: action_transfer_money
  - action: utter_transfer_complete
  - action: utter_goodbye

Here we define two conversation paths:

1. The user greets the bot, asks for their checking balance, and the bot fetches the balance and says goodbye.
2. The user asks to transfer $100, the bot executes the transfer and confirms completion.

Of course, this is just a tiny slice of a real banking assistant, but it illustrates the basic concepts. A full assistant would define many more intents, entities, and flows, along with error handling and clarification prompts.

The final piece is to implement the custom actions. Let‘s sketch them out in actions/actions.py:

from rasa_sdk import Action
from rasa_sdk.events import SlotSet

class ActionCheckBalance(Action):

    def name(self):
        return "action_check_balance"

    async def run(self, dispatcher, tracker, domain):
        account_type = tracker.get_slot("account_type")
        # TODO: Fetch real balance from database or API
        balance = 5000.00  

        dispatcher.utter_message(response="utter_balance", account_type=account_type, balance=f"${balance:.2f}")
        return []

class ActionTransferMoney(Action):

    def name(self):
        return "action_transfer_money"

    async def run(self, dispatcher, tracker, domain):
        amount = tracker.get_slot("amount")
        # TODO: Execute real money transfer 
        dispatcher.utter_message(response="utter_transfer_complete", amount=amount)
        return [SlotSet("amount", None)]

The custom actions use the Rasa SDK to interact with the conversation tracker and dispatcher. Here we simply return hard-coded balance values and transfer confirmation messages, but a real implementation would integrate with backend APIs and databases to fetch data and execute transactions.

Step 5: Train the NLU and Dialogue Models

With our training data defined, we can now train the Rasa models:

rasa train

This command trains both the NLU model to recognize intents and entities, and the dialogue model to predict the next action based on the conversation state and context.

Under the hood, the NLU model is using a transformer-based architecture (by default, DistilBERT) to classify intents and extract entities from user messages. The dialogue model is powered by a collection of machine-learning policies, including the TEDPolicy (Transformer Embedding Dialogue) and MemoizationPolicy.

The trained NLU and dialogue models are bundled together as a compressed tar.gz file in the models/ directory, with a timestamp in the filename.

Step 6: Talk to Your Assistant

Let‘s take our assistant for a spin! To load up an interactive chat session, run:

rasa shell

You should see a command prompt where you can type messages to your bot. Try out some of the sample conversations you defined in stories.yml, and see how the bot responds!

Your input ->  hi
Hi there! How can I help you today?

Your input ->  what‘s my checking account balance?
Your checking account balance is $5000.00.
Goodbye! Let me know if you need anything else.

Your input ->  transfer $100 from savings to checking
OK, I‘ve transferred 100 from your savings to your checking.
Goodbye! Let me know if you need anything else.

Remember, this is just the tip of the iceberg in terms of what you can build with Rasa. You can expand your assistant‘s capabilities by adding more intents, entities, and custom actions, integrating with external services and APIs, and deploying across multiple messaging channels.

Conclusion

In this guide, we‘ve walked through the process of setting up a powerful, cloud-based development environment for Rasa using AWS Sagemaker Studio Lab and code-server. With this setup, you have access to:

• Free, on-demand compute for training and developing your models
• A full VS Code IDE for writing code and configuration
• An isolated Conda environment for dependency management

We‘ve also sketched out the key components and workflow of a Rasa assistant:

1. Defining the assistant‘s capabilities in domain.yml
2. Providing training data for intents, entities, and dialogue flows in data/
3. Implementing custom actions in actions.py
4. Training the NLU and dialogue models with rasa train
5. Interacting with the assistant using rasa shell

Of course, we‘ve just scratched the surface of what you can do with Rasa. As an open-source framework with a vibrant developer community, Rasa provides a rich ecosystem of tools, integrations, and resources to help you build state-of-the-art conversational AI.

Some key resources to continue your learning:

• Rasa Documentation
• Rasa Community Forum
• Rasa Masterclass Video Course
• Rasa Blog

I hope this guide has been helpful in jumpstarting your Rasa development journey. Feel free to reach out if you have any questions or comments. Happy bot building!