System Design: Personalized Search with ML on AWS
As a continuation of our previous post, where we built a recommendation system with machine learning, we are now expanding our system functionality by adding a new text search for relevant videos using the power of machine learning.
This post will follow and mention some of the content in the “Building a Video Recommendation System Using Machine Learning on AWS” Post, so it is recommended to visit that post before jumping into this one.
By the end of this post, you will understand how to build and architect a personalized search system for videos using both text and video visual content.
Github code can be found in from here, keep in mind that the code in the repo can be tested locally in your own machine.
Keep in mind that for local demo data we do not have divesity in both video content and text content, thus you may see very similar results but with different scores
Table Of Content
∘ Table Of Content
· Support My Content
· Problem Framing
∘ Text Search System
∘ Visual Search System
· Data Exploration
· Feature Engineering & Data Preparation
∘ Statistical methods
∘ ML-based methods
· Video encoder
· Data Preparation
∘ Data Collection
∘ Preparing Textual Data
∘ Preparing Negative Text Query Data
∘ Preparing Video Visual Data
· Feature Store
· Algorithm Selection & Model Training
· Understanding the Training architecture
∘ Offline Metrics
∘ Extracting Training Data
· Model Serving and Deployment
∘ Online VS offline inference pipelines in recommenders
∘ Offline inference pipeline
∘ Online inference pipeline
· Inference options using SageMaker Endpoint
∘ Auto-Scaling The Endpoint
∘ Strategies to deploy new and updated models
· Model Monitoring
· Prediction & Data Collection
∘ Understanding the Collection of inference ground truth labels
· Machine Learning Pipelines
· Serving
∘ Indexing Pipelines Implementations
∘ Prediction Pipeline implementation
· Conclusion
· References
Problem Framing
The first step is to start a conversation with the client, understanding the requirements of the system and what are the expectation that the ML model will help us solve.
You can see the conversation we had with the client for the requirement gathering in the following diagram:
From the conversation we had, we can get high-level inputs to start with:
- Our system will expect text as input (User Query), and the output is a list of personalized video results for the user that match the query provided.
- User text query will be compared to both text-based and video-content-based.
- Results should be personalized to the user's preferences, meaning we show the user results based on their preferences.
- The text query language is English only for now.
Let us start by drawing a sketch of the model’s expected input and outputs:
As the client mentioned that our system should support both text-based search and video-content-based search, this means the system is built from two main parts:
- Text Engine: which performs text-based search, an example of this engine is ElasticSearch
- Visual Search Engine: which is able to understand video visual content and search for videos based on their visual content that matches the query text.
An overview of the design can be seen in the following diagram:
Let’s briefly discuss each engine component at a high-level.
Text Search System
Text search works when a user types in a text query: “How to develop an AI Agent”. Videos with the most similar titles, descriptions, or tags to the text query are shown as an output.
Text search databases support what is known as “The inverted index” which is a common technique for creating the text-based search component, allowing efficient full-text search in databases.
Since inverted indexes aren’t based on machine learning, there is no training cost. A popular search engine companies often use is Elasticsearch, which is a scalable search engine and document store.
Visual Search System
This component takes a text query as input and outputs a list of videos. The videos are ranked based on the similarity between the text query and the videos’ visual content.
The system contains three components:
- Video Encoder: generates an embedding vector from the video visual content.
- Text Encoder: generates an embedding vector from the text.
- Vector Database: A database that allows similarity search between two different embeddings in the same dimensional space.
Both Video encoder and text encoder will produce an embedding that lives in the same space dimension.
The similarity score between the video and the text is calculated using the dot product of their representations.
In order to rank videos that are visually and semantically similar to the text query, we compute the dot product between the text and each video in the embedding space, then rank the videos based on their similarity scores.
Data Exploration
Before diving into the system's internal data, we need to understand what our training dataset should look like for training. The data is pretty simple; it is a pair of a text query and a video path, or video_id which we can use to get the path of the visual content of the video. Lastly, we also have negative_text_queries which are N-1 negative samples of positive queries, as shown in the table below:
We cover each field generation in detail in the next few sections.
Having a look at the dataset we have used in the previous post about the recommendation system, going back to the Video dataset, which we expanded around it in this post, as seen:
This dataset can be a great start for our model development. We can use the title and video_caption fields as a way to start generating our model's positive text query dataset.
With those two fields, we can create two training rows for each video; both of these fields are captured from the user input when the user uploads a new video record, and we can force the user to fill in the video_caption field, or we can make it optional and filter the videos to only those that have this value filled.
Another dataset we can use is the interaction dataset:
We are already capturing the user queries that are running in the system, through the usage of the session_id field, we can track when a user types a query and follow it with a watch or a click, or even a like.
All these interactions can be a sign of positive interaction and can be used to generate a query from them.
QUERY ("Coding tutorial about AI") ──► SEARCH RESULTS ──► WATCH (video_id: 4, 8 sec) ──► LIKE (video_id: 4)A third way to generate the dataset is by human annotators, where we can provide humans to annotate videos with different queries. For each video, we can have a human annotate it with 5–6 different text queries (Positive or Negative).
Human annotation is usually expensive, and while it may be very accurate, it lack with speed and scale.
Last but not least, we can use AI to augment the dataset. We can use an LLM (like Claude or GPT) to generate descriptions from video frames or transcripts.
Regardless of what you choose, you will end up with a feature-rich dataset, which you can use for initial training and then start collecting actual labels from a real-time application pipeline.
Feature Engineering & Data Preparation
After understanding the data required for training and the data we have in internal systems, it is now time to start feature engineering.
Let us start with the text encoder; text queries are converted into embeddings by a text encoder, which converts text into a vector representation.
For example, if two sentences have similar meanings, their embeddings are more similar.
To build the text encoder, two broad categories are available:
- Statistical methods
- ML-based methods.
Let’s examine each.
Statistical methods
Those methods rely on statistics to convert a sentence into a feature vector. Two popular statistical methods are:
- Bag of Words (BoW): This method converts a sentence into a fixed-length vector. It models sentence word occurrences by creating a matrix with rows representing sentences, and columns representing word indices.
- Term Frequency Inverse Document Frequency (TF-IDF): This is a numerical statistic intended to reflect how important a word is to a document in a collection or corpus. TF-IDF creates the same sentence-word matrix as in BoW, but it normalizes the matrix based on the frequency of words.
In summary, statistical methods are usually fast. However, they do not capture the contextual meaning of sentences, and the representations are sparse. ML-based methods address those issues.
ML-based methods
In these methods, an ML model converts sentences into meaningful word embeddings so that the distance between two embeddings reflects the semantic similarity of the corresponding words. For example, if two words, such as “rich” and “wealth” are semantically similar, their embeddings are close in the embedding space.
The following figure shows a simple visualization of word embeddings in the 2D-2D embedding space. As you can see, similar words are grouped together.
There are three common ML-based approaches for transforming texts into embeddings:
- Embedding (lookup) layer: In this approach, an embedding layer is employed to map each ID to an embedding vector.
Employing an embedding layer is a simple and effective solution to convert sparse features, such as IDs, into a fixed-size embedding
- Word2vec: Word2vec is a family of related models used to produce word embeddings. These models use a shallow neural network architecture and utilize the co-occurrences of words in a local context to learn word embeddings. In particular, the model learns to predict a center word from its surrounding words during the training phase. After the training phase, the model is capable of converting words into meaningful embeddings.
There are two main models based on word2vec: Continuous Bag of Words (CBOW) and Skip-gram.
The following figure shows how CBOW works at a high level.
Even though word2vec and embedding layers are simple and effective, recent architectures based on Transformers have shown promising results.
- Transformer-based architectures: These models consider the context of the words in a sentence when converting them into embeddings. As opposed to word2vec models, they produce different embeddings for the same word depending on the context.
The following shows a Transformer-based model which takes a sentence — a set of words as input, and produces an embedding for each word.
Transformers are very powerful at understanding the context and producing meaningful embeddings. Several models, such as BERT, GPT3, and BLOOM, have demonstrated Transformers’ potential to perform a wide variety of Natural Language Processing (NLP) tasks.
In our case, we choose a Transformer-based architecture such as BERT as our text encoder.
Video encoder
We have two architectural options for encoding videos:
- Video-level models: process a whole video to create an embedding, as shown in the following figure. The model architecture is usually based on 3D convolutions or Transformers. Since the model processes the whole video, it is computationally expensive.
- Frame-level models: work differently. It is possible to extract the embedding from a video using a frame-level model by breaking it down into three steps:
— Preprocess a video and sample frames.
— Run the model on the sampled frames to create frame embeddings.
— Aggregate (e.g., average) frame embeddings to generate the video embedding.
Since this model works at the frame level, it is often faster and computationally less expensive. However, frame-level models are usually not able to understand the temporal aspects of the video, such as actions and motions.
In practice, frame-level models are preferred in many cases where a temporal understanding of the video is not crucial. Here, we employ a frame-level model such as ViT for two reasons:
- Improve the training and serving speed
- Reduce the number of computations
Data Preparation
We have chosen how to feature engineer each of the encoders that we have; for text, we used BERT, and for video, we chose a frame-level model called Vit.
It is time now to start processing the data that we have in the system to make it available for consumption by the end models and train those models.
Data Collection
Before jumping into the data preparation step, here is a high-level overview of how to collect data for training the model. We use a batch-based pipeline in this situation because it is easier to manage, plus the models do not need to collect new data in a streaming fashion; it is more than enough to grab new data in hourly-based or daily batches and re-train the model.
In a nutshell, the data is read by the two buckets (interactions and videos) and is extracted through A Glue ETL job or an EMR job. Once the data is processed, it is landed in a new bucket called the video-query bucket, in which we can start feature engineering the data.
Data can land in this bucket with the following format:
video_id, video_path, query_text
1 , s3://location, how to cook an ege
2 , s3://location, build an AI Agent with langchain
....When a new file or record lands in the Video-Query bucket, we can use the S3 events to trigger a job that will generate negative sampling for those positive queries or re-use the architecture we had shown previously by capturing DynamoDB CDC.
Preparing Textual Data
Let us start with the simplest method, which is the Text embedding. As mentioned, we used BERT in this case; for a deep dive into BERT check my other blog post on building a text classification model. We provide a high-level understanding of the BERT model in this post.
We can start running a data processing job on the video-query bucket; the result of this job will be saved into an offline SageMaker FeatureStore, where we can further use it to train the model.
In the processing job, we will download the BERT tokenizer, read data from the S3 bucket (new data or full data), and tokenize the text queries into BERT supported format, then we save the output results.
BERT tokenizer usually outputs the embedding into two fields input_ids and attention_mask
SageMaker FeatureStore allows us to store features and share them across the company. We can use the feature store offline with AWS Athena for point-in-time queries to travel in time with data.
Preparing Negative Text Query Data
As positive samples are generated from the data we own in the system, negative samples can be generated in two ways:
- Human-based generated: higher accuracy, higher cost, hard at scale.
- AI-generated: lower accuracy than humans, less cost, and works at scale.
For that, we obtin for AI-Generated data, and while it provides less accuracy, it works at scale, let us review a high-level architecture where we can use AI to obtain those values.
We used a similar approach to the positive queries in the preprocessing. For generations, when a new query is delivered into the S3 bucket, we use S3 events for firing a record into the SQS. A consumer lambda function will pick it up and call a Bedrock LLM to generate “X” negative queries.
Next, it stores them in a negative queries bucket, where a downstream processing job will run on this data and store them in an offline feature store as well.
Preparing Video Visual Data
Because we want to embed video at the frame(s) level, in that case, we need to follow a multiple-step approach ot handle frames:
Let's go over each step in the definition and then show a high-level architecture of the implementation:
- Decode Frames: Unpacks compressed video into raw frames
- Sample: Picks every Nth frame to reduce redundancy
- Resize: Standardizes all frames to the same dimensions
- Normalize: Scales pixels and adjusts for pretrained models
- frames.npy: Stores processed data for fast loading
At the final step the frames.npy is a numpy array of processed frames, the array would look similar to the following:
frames.npy = [
Frame 0: [[[ 0.2, 0.5, 0.3 ], [ 0.1, 0.4, 0.2 ], ...], ─┐
[[ 0.3, 0.6, 0.4 ], [ 0.2, 0.5, 0.3 ], ...], │
...] │
│
Frame 1: [[[ 0.4, 0.7, 0.5 ], [ 0.3, 0.6, 0.4 ], ...], │
[[ 0.5, 0.8, 0.6 ], [ 0.4, 0.7, 0.5 ], ...], │──► Feed to ViT
...] │
│
... │
│
Frame 7: [[[ 0.6, 0.9, 0.7 ], [ 0.5, 0.8, 0.6 ], ...], │
[[ 0.7, 1.0, 0.8 ], [ 0.6, 0.9, 0.7 ], ...], │
...] ─┘
]For each frame in the numpy array, we can pass it to the Vit model, which, in its turn, is going to patch-embed the video frames and aggregate them, and finally, we can store the output embedding in the SageMaker FeatureStore offline.
We used AWS Batch to run frame sampling on scale. Because we run this pipeline in batch mode, it is best to handle it on a distributed computing service such as AWS Batch.
For each video, we store the processed frames in an S3 bucket, then we that saved in a featureStore offline for querying it later on for training jobs.
While this option works very well for batch processing at scale, we need to calculate the video embedding on the fly when a new video is uploaded.
After newly uploading a video to the system, we need to make it available to be queried as soon as possible. We need to calculate its embedding and save it in a Vector Database. We can either re-use this pipeline to run it every X video upload or every X minute, or we can adopt a real-time solution with the following architecture:
Here is a quick overview of the Architecture:
- When a video is uploaded to S3 and written into a DynamoDB Table, we can capture the written record with the DynamoDB Streams.
- DynamoDB Stream will trigger a lambda function, which in its turn will send a record into the SQS queue.
- Later, depending on your choice and traffic pattern and scale, either a Lambda function or ECS Fargate will pick up the record from the queue, download the S3 file attached to the video, and generate
frames.npyfile from it, which will be inserted into an S3 bucket. - S3 event will trigger a new record in SQS, which a downstream consumer will pick up, and the consumer will call the SageMaker Video Encoder Model with the
frames.npyit got processed and received an embedding. - Once the embedding is generated, it will store it in a JSON format in the SageMaker Feature Store offline.
- The generated embedding will also be stored in a vector database for querying.
Feature Store
The last part to cover in this section is the feature store data structure we want to use. We want to understand how our data is going to be stored in the feature store, as we have seen this being mentioned coupel of times.
- The
video_fgstoresvideo_idand the path to the generated frames of the video (one record per video) - The
query_samples: Because each video can have multiple positive samples, we store each sample with thevideo_id, and we store each text BERT tokenized in theinput_ids_jsonandattention_mask_jsonThe label field is used to know whether the review is a positive review (1) or a negative review (0). Finally, thepair_idis critical that we connect the negative pairs with the positive ones.
When a video is uploaded and samples are generated, we insert it into the video_fg When a positive query is generated, we upload it to the query_samples_fg and when a negative one is generated, we upload it to the same feature group (query samples) as well.
Algorithm Selection & Model Training
As we already know what algorithm we want to choose to train the two encoder models, it is time to start adapting those algorithms and set up the training pipelines on AWS.
As a quick reminder, AWS provides us with the following ways to train our ML model:
In the case where we want to deploy a simple ML model, it is recommended to either use pre-trained models or the built-in algorithms; however, AWS also supports pre-made images such as TensorFlow and Pytorch, which makes a great fit for building our ML models.
Understanding the Training architecture
To train the text encoder and video encoder, we use a contrastive learning approach. An explanation of how to compute the loss during model training is shown in the following diagram:
- The video encoder will accept a
frame.npy(a numpy array of frames) - Text encoder will accept one positive and N-1 negative samples, N should be between 5 and 10, which means for a single query text, we need to generate 5–10 negative queries (we covered this in a previous section)
As the diagram suggests, the architecture is divided into two towers:
- The Text Tower (Query Encoder): Encodes the search query or caption text using a pretrained BERT model. The raw text (e.g., “A dog playing in the park”) is first tokenized into
input_idsandattention_masktensors using the BERT tokenizer. - The Video Tower (Item Encoder): Encodes the visual content of the video using a pretrained ViT (Vision Transformer) model. The input consists of 8 uniformly sampled frames from the video, each resized to 224×224 pixels and normalized.
Both the Text Tower and Video Tower follow similar steps in embedding their respective inputs:
- Feature extraction with pretrained backbones:
— In the Text Tower, the tokenized caption is passed through BERT-Base, and the [CLS] token embedding (768-dim) is extracted to represent the entire text.
— In the Video Tower, each of the 8 frames is passed through ViT-Base to produce per-frame features (768-dim each). These are then combined using a learned attention mechanism that weights each frame’s importance, producing a single 768-dim video representation. - Projection to shared embedding space: Both towers use a feedforward neural network (Dense → GELU → Dense) to project their 768-dim features into a shared 256-dimensional embedding space. The outputs are L2-normalized to enable efficient similarity computation via dot product.
The choice of 256 dimensions is a practical trade-off.
┌─────────────────────────────────────────────────────────────────────────────┐
│ EMBEDDING DIMENSION TRADE-OFFS │
├─────────────────────────────────────────────────────────────────────────────┤
│ │
│ Too Small (16-64) Just Right (128-512) Too Large (768+) │
│ ───────────────── ──────────────────── ──────────────── |
│ ✅ Fast search ✅ Good accuracy ✅ Best accuracy │
│ ✅ Low storage ✅ Reasonable speed ❌ Slow search │
│ ❌ Poor accuracy ✅ Moderate storage ❌ High storage │
│ │
└─────────────────────────────────────────────────────────────────────────────┘Depending on your application requirements, feel free to adjust the Dimension value to suit your use case. Some applications will require lower accuracy while preferring Faster search, while others will look for more accurate results than speed.
Here is the math behind choosing the best dimensions:
| Dimension | Storage (1M videos) | Search Speed | Accuracy |
| -----------------: | ------------------: | ------------ | --------- |
| 64 | 256 MB | Very Fast | Lower |
| 128 | 512 MB | Fast | Good |
| ⭐ 256 | 1 GB | Good | Very Good |
| 512 | 2 GB | Slower | Excellent |
| 768 (raw BERT/ViT) | 3 GB | Slow | Best |The two encoders independently generate embeddings that live in the same dimensional space. By optimizing a similarity function (like a dot product) between these embeddings, the model learns to bring text closer to the videos they are likely to match visually.
Offline Metrics
Here are some offline metrics that are typically used in search systems. Let’s examine which are the most relevant.
- Precision@k: precision@k = number of relevant items among the top k items in the ranked list k precision@k =k number of relevant items among the top k items in the ranked list.
- Recall@k. This measures the ratio between the number of relevant videos in the search results and the total number of relevant videos.
- Mean Reciprocal Rank (MRR). This metric measures the quality of the model by averaging the rank of the first relevant item in each search result. The formula is:
Extracting Training Data
Let us go into a deep dive on how the training data should be extracted:
Looking into the diagram, let us review the extraction of training data:
- Using AWS Athena and the offline feature store we have designed, we query the offline feature store to get the latest “X” hour or month data.
- Queried data is exported into an S3 bucket, where we store it with a version number. For example, this is the 3rd time we train our model, we put the data into the following format
/visual-search/full/v3At this level, data is not split into training, validation, or testing. - We run a SageMaker processing job to split the queried data into training and validation datasets, and we store them in the same format as the previous step. In this step, we convert our data into a TfRecords which is optimized for training.
- Using the split data, we run the SageMaker training job by pointing it at the data we want to train on.
- Training will run, finish, and save the model into a model registry versioned with the model version, and data version (path to the data)
Model Serving and Deployment
Once the models are trained, the next step is making them accessible to end users and applications. Finally, it is time to deploy these models into production.
Before going into the implementation details, we want to explain the serving strategy of our inference pipelines. We have one offline and one online inference pipeline.
Let’s understand the difference between the two in our video search.
Online VS offline inference pipelines in recommenders
The inference pipeline is split into two main processes to optimize for real-time recommendations.
- The offline pipeline runs in batch mode, optimized for high throughput. It embeds all the video frames from our database using the video encoder.
The offline pipeline runs once to backfill our video collection. It should then run again whenever a new video is added to our collection or the model has been retrained (which changes our embedding space). - The online inference pipeline is deployed as a real-time service optimized for low latency. It will run on each client request. We also deploy the video encoder as an online, real-time service for serving new videos that are uploaded.
Let’s zoom in on each pipeline.
Offline inference pipeline
The offline pipeline loads the video encoder from the SageMaker model registry and a reference to the video feature group stored in the SageMaker feature store.
Leveraging the feature store, it uses the video_path to run the video encoder on all the available videos in our system.
Ultimately, it saves the video embeddings into a vector database that supports a vector index for semantic search between the videos and the text query.
Why do we use offline pipelines? Handling all the system videos with an online endpoint will require heavy computational processes; it will cost too much to run the online service and pass each video to it both computationally and in terms of network.
So, to handle new videos, we have two ways to do that:
- Deploy a lightweight online version of the model
- Run the offline pipeline when X new videos arrive or every X minutes
Online inference pipeline
In the online inference pipeline, both of our models will serve online requests. The reason for this is that when a user types in text in the input, we need to call the Text Encoder, then perform a similarity search on the vector database, and when a new video is added to the system, the video has to be passed across the Video Encoder and then saved into a vector database.
Let us review a high-level design of the visual search system to understand why we need the endpoints
Given the flow of the data we have in the system, we see that we need two different endpoints one to deploy the text encoder (powered by BERT) and the other one for the video encoder (powered by Vit).
If you are optimizing for cost, and not real-time you can drop the online video encoder endpoint, and replace the flow with an offline one that runs on the new videos every X hour
Inference options using SageMaker Endpoint
With SageMaker, we can deploy our model into an endpoint that we can use and get inference from.
SageMaker gives us multiple options for Real-Time predictions:
- Real-Time Inference: Real-time inference is ideal for inference workloads where you have real-time, interactive, low-latency requirements.
- Serverless Endpoints: Serverless Inference is a purpose-built inference option that enables you to deploy and scale ML models without configuring or managing any of the underlying infrastructure.
- Async Endpoints: Amazon SageMaker Asynchronous Inference is a capability in SageMaker AI that queues incoming requests and processes them asynchronously. This option is ideal for requests with large payload sizes (up to 1GB), long processing times (up to one hour), and near real-time latency requirements.
Choosing the best inference option depends on your application traffic patterns and the features you want to use from the endpoints. In our case, we opt for Serverless inference because we do not have newly uploaded videos all the time, plus the usage pattern of the search system is not balanced and can not be determined. The best choice for this case is a serverless endpoint for both encoders (keep in mind that we can use real-time inference for the text-encoder and serverless-inference for the video-encoder, but again, this solely depends on your traffic pattern.
With serverless inference, we do not worry about adding an extra layer of management; SageMaker takes care of the provisioning, management, and scaling of that endpoint.
Auto-Scaling The Endpoint
Amazon SageMaker AI automatically scales in or out on-demand serverless endpoints. For serverless endpoints with “Provisioned Concurrency” you can use Application Auto Scaling to scale up or down the Provisioned Concurrency based on your traffic profile, thus optimizing costs.
Strategies to deploy new and updated models
When working with serverless endpoints, you’ll quickly notice they don’t support blue/green deployments out of the box.
So, how do we handle rolling out new model versions safely?
One approach that works well is spinning up two separate endpoints, each with its own configuration. Your application servers then use AWS AppConfig to determine which endpoint to call — and more importantly, how much traffic each one should receive.
The idea is simple: start by sending just a small percentage of traffic to the new endpoint, then gradually increase it over time while keeping a close eye on your metrics. If something looks off, you can quickly roll back. The diagram above illustrates how this configuration-driven deployment flow works in practice.
Model Monitoring
Once our model is live, it should be monitored; monitoring involves capturing the performance of our model and the data we are receiving.
Usually, you need to focus on four metrics:
- Model Performance ~ different between model prediction on training vs online data.
- Data Quality ~ quality of the data in production vs training.
- Model Bias ~ model biasing a specific attribute value rather than another in the prediction.
- Model Attribution Shift ~ model favouring a specific attribute over the other for the prediction.
Model Performance and Data quality must be monitored in all cases, while bias and attribution shift can be avoided in our case, because our model only uses text field for the query encoder, and a frames.npy array for the video encoder.
For different cases where we have an attribute value that may contain a categorical value, we may need to monitor the Model Bias, and when we have multiple inference fields, we have to monitor the attribution shift.
Following is a high-level architecture for monitoring both model performance and data quality. The only differences between them are:
- For model performance, we need to collect predictions (model inputs + outputs) and ground truth labels.
- For model data quality, you only need to collect the model predictions (model inputs + outputs); no need for ground truth labels.
Monitoring Jobs can be run on a schedule, for example, every two hours.
A key important note is how you want to capture your endpoint inputs and outputs, which depends heavily on what type of endpoint you are using.
- If you are using Real-Time Endpoint, you can turn on the “Data Capture” features, which allow you to capture all your endpoint requests and output predictions.
- If you are using a serverless inference (Such in our case), you will have to build the logic of collecting data on your own. Here is a simplified architecture example of how you can do this:
One more part of monitoring that we need to measure is the online metrics that are related to the visual search functionality, for which we can measure some online metrics, such as:
- Click-through rate (CTR)
- Video completion rate
- Total watch time of search results
- Abandoned interactions (Interactions that start with a query and did not end with a watch/like)
Due to the extensive and repetitive implementation, I would recommend you visit the following two posts for more information and implementations:
Prediction & Data Collection
Our model is now serving users, predictions are made, and monitoring is set in place for any drift changes. It is now time to set up a plan on how we can capture those predictions and how to collect the ground truth labels for each prediction.
In the previous stage, we have enabled data capture flow by developing our own custom pipeline, and we have been able to capture the data input and output from the model.
However, we are still missing one important answer: how can we collect the inference ground truth labels and merge them with model predictions?
Understanding the Collection of inference ground truth labels
When we talk about collecting inference ground truth labels, we are referring to the process of knowing if the model prediction was correct or not, and thus, we need to keep track of the model output and collect the user feedback through the application.
Following is a high-level general architecture of how to collect the inference ground truth and merge it with our results:
So to collect the ground truth labels, we can keep track of the session the user has initiated a text query with, then track that session if the user ended with a like/watch/interaction with in the same session, we can consider that to be a positive labels, sessions that contain a text query that does not endup with one of the previous interaction will be considered as a negative label.
To collect Data from the user text queries:
- Any text query → video list results → watch/like can be considered as a positive pair.
- Any text query → video list results → ignore/dislikes, etc., can be considered as a negative pair.
We can collect the positive query → video pairs and use them for training. Similarly, we can do the same for negative pairs, say for example we return user 10 videos and none of them have been watched or liked by the user, we can create a negative query → video_id pairs for training as well.
To improve accuracy over the negative paris generated, we can send those into a human to be labeled, we can use SageMaker GroundTruth Streaming for manual labeling at scale.
For implementation details, make sure to read this blog post:
Machine Learning Pipelines
A Machine Learning Pipeline is a systematic workflow designed to automate the process of building, training, and deploying ML models.
A generic pipeline will start by preprocessing data, loading the training and validation datasets with versioning, evaluating the new model with the right metrics, and deploying the new model into production.
Let us go through a high-level design for our machine learning pipeline in this case:
In general, we will use AWS Step Function to run the full ML pipeline. Step Function can be decomposed into more step functions, each for a specific job or stage:
- Preprocessing Data Jobs: As data is continuously ingested into different S3 buckets (Raw data), we want to run a process job on this data with a SageMaker Processing Job. This job will focus on generating two different kinds of data:
1. Tokenizing the Text data from the video/query dataset into BERT supported format
2. Generatingframes.npyfor each newly uploaded video if handled in batch mode
The processing jobs can run in sequential or in parallel, depending on the data structure and system needs. - Split Data With Versioning: We run another processing job on SageMaker for splitting the data into training/validation/testing, and we insert that data into an S3 bucket with a specific format (TFRecords), saved in the following structure such as
/interaction/training/v3where thev3Is the training versioning number - Training Job: training job code lives in Github for example, when the pipeline run we download the code and point the training job into an S3 input with the right training data version. Depending on how we trigger our pipeline, we either do continuous learning on new data only, or a full training on all the data we have.
- Model Registry: With the model register, we save the output models from the training job into it, and with each model, we add metadata such as model version and training data version as well.
- Deployment Phase: Depending on the model we have (Text Query Encoder and the Video Encoder~ Online Deployment) or (Video Encoder for offline embedding), we deploy our new model into an endpoint and set up monitoring for it. We have covered the new deployment strategy in the previous section.
The Pipeline can be triggered in one of two cases:
- Full Training: We load all the data we have currently in the SageMaker Offline Store and train our model on it.
- Continuous Training: Only load the new data that has been ingested since the last training Job.
Serving
Let us focus now on how our overall application will work and interact with each other.
At serving time, the system will return textual and visual results that are similar to the user query provided.
In the early sections of this post, we mentioned that we are going to use ElasticSearch for Text-based search, and we have developed our model (Text and video encoders) for visual search.
We can split the serving into three different paths:
- Video Indexing Pipeline: Responsible For Indexing Video Visual Content.
- Text Indexing Pipeline: Responsible for filling up the text database (ElasticSearch) with the newly ingested data.
- Prediction Pipeline: Takes in the user text query and calls both visual and text search systems, and returns a list of videos to the user that are ranked by user preferences.
Let us start with the implementation details on AWS. We start with the indexing pipelines (Text + Visual)
Indexing Pipelines Implementations
Assuming our video database is a DynamoDB database, we start by enabling DynamoDB Streams to listen for newly added videos records, as shown in the diagram
Let us go over the architecture real quick:
- Once a new video is uploaded to S3 and inserted in DynamoDB, a DynamoDB stream record will trigger a downstream Lambda consumer.
- Lambda will read the record, and because the video needs to run against different indexing pipelines, it uses SNS as a fanout service to publish a new record to each pipeline's indexing queue.
- For the text pipeline queue, another lambda picks up the data and calls the auto-tagger service, which can be either another ML model or an LLM model deployed through BedRock. Once the auto-tags are generated, we insert the new record in the OpenSearch Database for text search.
- For the video visual indexing pipeline, Lambda will download the original video content from the S3 bucket, extract frames from it, and generate a
frames.npythat is passed to a video encoder deployed on SageMaker Endpoint. Once the Embedding is generated from the model, we can insert it into the Vector database for visual content similarity.
Because video processing can be expensive in resources and takes longer than 15 minutes we can switch the usage of AWS lambda in the visual content indexing pipeline with an ECS or EC2 instances.
Prediction Pipeline implementation
As both text and visual search databases are now filled with the new video data, we can start querying from those databases. When a user sends a query, it goes through the “Prediction Pipeline” service, which in turn runs in parallel to call both search engines we have built.
The following diagram shows the implementation architecture
Let us go over the architecture real quick:
- When a user types a query and hits enter on their keyboard, we send a GET request to the API Gateway with the text the user has typed in.
- The lambda function will run the request in parallel in two engines (Text Search Engine and Visual Search Engine)
- Text Search Engine will use an AWS OpenSearch client and perform a search with the user query.
- Visual Search Engine will call the Text Encoder with the query text, the model will use the
input_handlerto transfer raw text into BERT Tokenized input, which the model can understand, then it will call the Text Encoder Model to generate an embedding for the query. - The generated Embedding will be used for similarity search in a Vector Database.
- As both engines finish processing, a fusion layer will be invoked with both engines' results. Fusion will be an in-memory component that will filter down the results, for example, filter duplicated videos, and apply some business logic that we can allow from a filtering UI in the search page, etc.
- Once we have our final candidates, we send them to the ranking service, which is the same as the one we have developed in the previous blog post (Recommendation System), which will rank those videos depending on user preferences.
- Ranked results will be paginated and returned to the user sorted according to their preferences.
Conclusion
As your system evolves and you are required to ship new features, you need to think of ML as more than just a regression or a classification model; it can be used for major cases such as recommendation, text search, audio search, etc.
While we have shown you how to build a visual search, the system principles can be applied to add more engines to your search system. For example, following a similar architecture, you can add a third search engine that does a search for videos by the video’s audio content.
If you have any questions, feel free to drop them in the comment section.
References
- ByteByteGo Machine Learning Course
- AWS SageMaker AI Deep Dive Part I (Introduction, Data Labeling, Preparation, Processing, and Feature Store)
- AWS SageMaker AI Deep Dive Part II (Model Training, Hyperparameter Tuning, Debugging, and Model Registry)
- AWS SageMaker AI Deep Dive Part III (Model Inference, Endpoints, and Deployment)
- AWS SageMaker AI Deep Dive Part IV (Model and Data Quality Monitoring, MLOps)
- The Only Machine Learning Framework You Will Need
