Multiclass Text Classification on AWS SageMaker

Executive Summary

This project demonstrates a complete end-to-end machine learning pipeline for text classification using modern transformer models. We leverage DistilBERT (a lightweight BERT variant) to classify news articles into categories, with the entire workflow hosted on AWS SageMaker.

Architecture Overview

---

Key Technologies

Component	Technology	Purpose
Model	DistilBERT	Lightweight transformer for text classification
Dataset	News Aggregator (UCI)	400k+ news headlines with 4 categories
Training	SageMaker Training Job	Managed distributed training
Serving	SageMaker Endpoint	Real-time inference with auto-scaling
Inference	AWS Lambda	Serverless request handling
Libraries	Hugging Face Transformers	Model implementation and tokenization

---

Project Workflow

Phase 1: Exploratory Data Analysis

Before training, we conduct thorough EDA to understand the dataset characteristics, class distribution, and text patterns.

Loading notebook...

---

Phase 2: Model Training

We fine-tune DistilBERT on the news classification task using Hugging Face’s Trainer API integrated with SageMaker.

Loading notebook...

---

Phase 3: Deployment & Inference

The trained model is deployed to a SageMaker endpoint with Lambda integration for scalable serverless inference.

Loading notebook...

---

Lambda Inference Handler

The serverless inference is handled by an AWS Lambda function that invokes the SageMaker endpoint:

import boto3
import json

runtime = boto3.client("sagemaker-runtime")
ENDPOINT_NAME = "multiclass-text-classification-endpoint"

def lambda_handler(event, context):
    payload = json.loads(event["body"])
    
    response = runtime.invoke_endpoint(
        EndpointName=ENDPOINT_NAME,
        ContentType="application/json",
# Deployment code snippet
predictor = huggingface_model.deploy(
    initial_instance_count = 1,
    instance_type = "ml.m5.xlarge",
    endpoint_name = "multiclass-text-classification-endpoint-final-v1"
)

Key Learnings

1. SageMaker + Hugging Face Integration: The HuggingFace estimator simplifies training transformer models on managed infrastructure.

2. Endpoint Scaling: Auto-scaling policies ensure cost-effective deployment during variable traffic.

3. Serverless Architecture: Lambda + SageMaker endpoints provide a production-ready, pay-per-use inference solution.

4. Model Optimization: DistilBERT provides 60% of BERT’s size with 97% of its performance—ideal for cost-sensitive deployments.