Building a platform for GenAI

The Adventure Begins

It was January 2023, and I was discussing with my manager, Rick, what we wanted the Data Science (DS) Platform team to work on in the next quarter. He mentioned that I should block out some of my time to look into recent developments in natural language processing (NLP), since the field had advanced a lot since our DS team last deployed a major NLP project in 2021. I didn’t know it then, but I was about to embark on what has been the most exciting and rewarding project of my career so far.

The GenAI boom was just starting to heat up: ChatGPT had launched in November 2022, and GPT-4 had yet to be released. The time was ripe to start exploring the possibilities of what these new technologies could do for our company.

Early Experiments

I finally got started on our NLP 2.0 exploration in March 2023. This was excellent timing: GPT-4 was released on March 14. During that week, I tested GPT-4 on several tasks that my coworker had previously tested with GPT-3.5. While GPT-3.5 had failed these tasks just the week before, GPT-4 successfully completed them. It represented a major breakthrough in the practical possibilities of product applications that just weren’t feasible before. Looking back, I’m still amazed at what GPT-4 could do on tasks I was experimenting with for the first time.

In April I gave a presentation to our DS team that I subtitled: “What have I been doing for the past month?” I gave an overview of the ideas that we had been brainstorming along with some Product Managers, focusing in on three main ideas that we thought were most promising: 

1. Combining large language models (LLMs) with ongoing work on indexing the Knowledge Base that is referenced when members contact us for support
2. Intent and Action Classification: improving our text classification using LLMs
3. Automatic generation of notes from calls and chats: speeding up manual documentation that care coordinators were writing after each encounter with a member.

The Birth of Wordsmith

By May, we had done enough initial validation of these different ideas that we knew we wanted to start investing in some kind of platform for LLMs and GenAI, which I named Wordsmith.

At a high-level, Wordsmith was meant to do anything you wanted with text data, and to be so easy to use that other teams could just treat it as a black box. This is the first diagram from my original plan, and one I still frequently reference:

Inside that black box, there would be several parts. In the plan, I laid out the five main components of the platform:

1. NLP Proxy Service and Client Library
2. Prompt Tooling
3. Training Library
4. Evaluation Library
5. Model Inference/Model Serving System

Development of the platform has generally proceeded along this breakdown, so I’ll go over each component and the progress that the team has made so far:

Proxy Service and Client Library

Given the strong abilities of API-based foundation models, as well as the growing number of models available from different providers, this has been the heart of our LLM platform thus far.

When we were starting out, there were a few libraries that provided a single interface to make LLM calls and would send your request to different providers from the client side. However, I made the decision early on that the logic of routing between model providers would be handled on the server. Instead of the request going from a Data Scientist’s local environment directly to Google, or OpenAI, all LLM requests at Included Health pass through a single internal service called wordsmith-proxy, which then routes them to the appropriate provider.

Proxying on the server side has several practical benefits: the first is that users are insulated from churn as we work on routing logic and added providers. They don’t have to upgrade the client package to get each new change. It also makes provider credential management easier: instead of trying to figure out how to securely distribute separate credentials to each user, we could set up a single OpenAI API key, a single Google service account credential, etc that belongs to the proxy service.

How we actually implemented this: on the server, we followed the spec for the OpenAI API, which they publish in OpenAPI format. I figured that OpenAI probably had already written a better API for communicating with LLMs than I could come up with. Because we’re serving the same API as OpenAI, our Data Scientists and other internal users can just use the OpenAI Python SDK, and point the base URL at our internal endpoint.

When the request gets to our server, we translate it from OpenAI format into a request to one of four different providers: OpenAI, Google VertexAI, AWS Bedrock, and internal (wordsmith-serving). In the case of OpenAI, obviously, the translation is just a pass-through function: no modifications are needed.

For our Data Scientists, this architecture means that they can switch between model providers without needing to write any new code, just by changing the model name from a string like “azure:gpt-4-turbo” to “google:gemini-pro”. It also makes us cross-language out of the box: our engineering team mostly writes in Go, and recently integrated with Wordsmith Proxy using the Go OpenAI library.

I believe that making this choice for our platform has paid off. Since then, open-source LLM server tools have popped up that use the OpenAI schema as a common interface. If I were starting today, I would strongly consider using something like LiteLLM. Additionally, more client tools like LlamaIndex are supporting connections to “OpenAI-like” APIs out of the box.

Model Inference/Model Serving System

For online serving, we deployed MLServer internally, and use the HuggingFace runtime to serve models. This service is connected to our internal deployment of MLFlow, and it can download model artifacts from there. For batch inference, my teammate Amogh Nalwaya built out our wordsmith-inference system, which integrates with our data warehouse to allow users to launch jobs that apply LLMs over text data stored in the tables.

Training Library

At the outset, our Kubernetes infrastructure didn’t support GPUs, so the first task was to work with our Infrastructure and Data Infra teams to get that working. We use Karpenter for provisioning nodes in our compute cluster, so we worked to set up a new Provisioner to create GPU nodes. I also deployed the nvidia-device-plugin so that we could request GPUs as Kubernetes resources.

On the software side, we’re using HuggingFace transformers. Most of what I worked on for training was applying the basic examples from HuggingFace to our particular datasets, and experimenting with some more advanced techniques like Lora.

Evaluation Library

The wordsmith-evaluation library works as a wrapper around HuggingFace evaluate. You can use it with any of the existing metrics in the evaluate library, and we’ve also integrated some open-source metrics that aren’t in the library, as well as custom metrics that we’ve developed in house.

Prompt Tooling

Jack Sullivan led development of a Python library to abstract away boilerplate from prompting. For example, it makes it easy for users to generate prompts from string templates, which is a common use case for LLM applications. It also exposes a simple interface for users to provide examples for few-shot prompting, and then within the library can select the desired number of examples using semantic similarity.