Using API Gateway, Lambda, SageMaker and DynamoDB to build a categorisation service in AWS

Posted by on January 19, 2024

I’ve talked previously about the value of combining rules-based and machine learning approaches to categorisation. In short, rules-based approaches make it easy to do customer-level personalisation that complements a machine learning model trained to find patterns across customers.

In this post I’ll talk about how we used AWS to build an expense categorisation service that combines machine learning with a rules-based approach. This service forms part of the Smart Capture feature of our Ruby on Rails application, where we automatically extract data from receipts and provide suggested accounting categories.

The model

We began creating our expense categorisation model with a period of experimentation to decide the model architecture, such as the input features and algorithm. We knew that we would ultimately be hosting our model on a real-time SageMaker endpoint so we made sure to pick an algorithm that would meet our latency requirements as well as focusing on features that would be available for our app to pass to the model. 

Once our model architecture had been determined we set up automations for generating training data, model training and deployment using GitHub Actions. We then had an endpoint that we could send requests to and get back predictions in real-time. 

The similar expense lookup

The other half of the expense categorisation service is a rules-based similar expense lookup1. When presented with a new receipt to categorise the system looks for previous similar expenses from that business. If a match is found then the category used for the similar expense is returned as a prediction.

At its core, this lookup involves querying a data store. We take the information about a new receipt, query a store of expenses and see whether there’s a match for a given business. We achieve this using Lambda and DynamoDB in AWS. We have a DynamoDB table containing past expenses for businesses that we can query using a Lambda function to find expenses that are similar to the receipt being processed. DynamoDB is well suited for our data store because it’s simple to use, cheap to run and fast. It had also already been used successfully for other projects by different teams.

Putting it all together

With our two methods for categorising expenses – the model and the lookup – we needed a way to serve them up to our app via an API. As part of processing a receipt we needed to send requests to an endpoint and get back predicted categories for if/when the receipt is converted to an expense.

API Gateway, which had already been used for other projects, felt like a natural fit for creating the API, particularly the integration between API Gateway and Lambda. Requests go to API Gateway which hands them on to a Lambda function to do some stuff and return predictions back to the app. We can have our Lambda function do whatever we like so long as we return data in a consistent format.

So what exactly does our Lambda function do?  First of all, it talks to a SageMaker endpoint to get predictions from the model using data from the request to API Gateway. Additionally, it queries the DynamoDB lookup table to find similar expenses.

The diagram below shows how things work when generating predictions for a receipt:

  1. The app sends a request to our API which includes the necessary data
  2. API Gateway hands off this request to a Lambda function
  3. The Lambda function makes a request to SageMaker and reads from the DynamoDB table then combines the responses together
  4. API Gateway passes the response from the Lambda function back to the app

Returning predictions for a given input is a big part of our expense categorisation service, but it’s not the whole story. As I’ve mentioned, we have a expense table that underpins looking up similar expenses. That whole process would be pretty poor if we didn’t also keep our lookup table updated.

Keeping the similar expense lookup up to date

If a customer is creating expenses from multiple receipts then we want to learn from how they categorise them as quickly as possible. Imagine I’ve got ten ‘coffee shop’ receipts in my wallet. If I create an expense for the first one it would be great if that could be used as a similar expense for categorising the next nine. To do this we take advantage of events published by our app. 

Our Ruby on Rails app publishes events when various different things happen, from creating an expense to processing a bank transaction. These events are processed by our event system in AWS before landing in a data lake.  

When an expense is created in the app we use the published event to update our DynamoDB lookup table. Here we take advantage of the integration between Lambda and EventBridge to trigger Lambda functions based on particular events. We do similar updates if an expense is changed or deleted. This whole event-driven process keeps our lookup table updated such that we can rapidly learn from customers’ actions to give the best possible categories.

We can expand the previous diagram to include this event-driven process alongside how we generate predictions: the app publishes events which are processed in AWS and trigger Lambda functions to update our DynamoDB table.

Summary

Here we’ve shown how native AWS services can be used to build a categorisation service behind a simple API interface. API Gateway allows us to decouple the details of our service from the interface used by our app. This means we could change the backend however we like without needing to update the app. The flexibility of Lambda allows us to implement rules-based logic alongside using a machine learning model in SageMaker, while DynamoDB provides a low-latency storage solution.

  1. We could also describe this as a type of symbolic AI, though that feels a little grandiose given its simplicity! ↩︎

Thanks to Owen Turner for feedback and suggestions on the structure of this post

Leave a reply

Your email address will not be published. Required fields are marked *