Refine BGE embedding model with synthetic data from Amazon Bedrock in 80 chars

Are you struggling to find high-quality data to improve your machine learning (ML) models? Synthetic data generation could be the answer, especially when real-world data is limited or sensitive. For instance, creating a medical search engine without access to real user queries and medical documents is a challenge due to privacy concerns. Fortunately, synthetic data generation techniques can help create realistic query-document pairs, enabling accurate model training while respecting user privacy.

In this blog post, we will show you how to utilize Amazon Bedrock to generate synthetic data, fine-tune a BAAI General Embeddings (BGE) model, and deploy it using Amazon SageMaker.

Amazon Bedrock is a fully managed service that offers various high-performing foundation models from top AI companies, including AI21 Labs, Cohere, and Amazon. It provides a wide range of capabilities to build generative AI applications with a focus on security, privacy, and responsible AI.

You can access the complete code for this post on our GitHub repository.

Solution Overview

BGE refers to Beijing Academy of Artificial Intelligence (BAAI) General Embeddings, a family of embedding models designed to produce high-quality embeddings from text data. The BGE models come in three sizes:

• bge-large-en-v1.5: 1.34 GB, 1,024 embedding dimensions
• bge-base-en-v1.5: 0.44 GB, 768 embedding dimensions
• bge-small-en-v1.5: 0.13 GB, 384 embedding dimensions

The BGE model operates as a bi-encoder architecture for comparing two pieces of text by processing them in parallel to obtain their embeddings.

Synthetic data generation can enhance model performance by providing high-quality training data without traditional constraints. This post will guide you through the process of generating synthetic data using Amazon Bedrock, fine-tuning a BGE model, evaluating its performance, and deploying it using SageMaker.

The key steps are:

1. Set up an Amazon SageMaker Studio environment with the necessary IAM policies.
2. Open SageMaker Studio.
3. Create a Conda environment for dependencies.
4. Generate synthetic data using Meta Llama 3 on Amazon Bedrock.
5. Fine-tune the BGE embedding model with the generated data.
6. Merge the model weights.
7. Test the model locally.
8. Evaluate and compare the fine-tuned model.
9. Deploy the model using SageMaker and Hugging Face Text Embeddings Inference (TEI).
10. Test the deployed model.

Prerequisites

First-time users require an AWS account and an IAM user role with specific permission policies, including AmazonSageMakerFullAccess and custom IAM policies for necessary permissions.