Intent Extraction Service Using DistilBERT

1. Overview

This document describes the design and implementation of an intent extraction service for a conversational AI application.

The goal is to reliably classify user input into one of three intent categories—consultation & Q&A, ideation & brainstorming, and planning & scheduling—using a fine-tuned DistilBERT-based model.

The service will be deployed as a REST microservice within a Kubernetes environment for scalability and high availability.

2. Requirements

Functional Requirements

• Intent Classification:?The system must classify incoming user queries into predefined intents: Consultation & Q&A, Ideation & Brainstorming, Planning & Scheduling
• Real-Time Processing:?The classification must be performed with low latency to ensure a responsive user experience.
• RESTful API:?The service should expose a REST API endpoint for receiving requests and returning intent classifications.

Non-Functional Requirements

• Scalability:?The microservice must support horizontal scaling via Kubernetes.
• Maintainability:?The design should facilitate easy updates and model retraining.
• Cost Efficiency:?The solution should minimize computational overhead without sacrificing accuracy.
• Security:?Secure communication between the client and the service must be maintained.

3. System Architecture

High-Level Architecture

1. Client Application:?Sends user queries via HTTPS to the intent extraction service.
2. API Gateway:?Routes incoming requests to the appropriate microservice.
3. Intent Extraction Microservice: 1) Preprocessing Module:?Normalizes and tokenizes the incoming text. 2) DistilBERT-based Classifier:?Processes the tokenized input and outputs an intent label., and 3) Postprocessing Module:?Formats the response and sends it back to the client.
4. Logging & Monitoring:?Captures performance metrics and error logs.
5. Kubernetes Cluster:?Orchestrates the deployment and scaling of the microservice.

Data Flow Diagram

User Query -> API Gateway -> [Preprocessing] -> [DistilBERT-based Classifier] -> [Postprocessing] -> API Gateway -> User Response        

4. Model Selection and Architecture

Model Choice: DistilBERT Model

Why BERT?

• Contextual Understanding: BERT-based models, including lighter variants, excel at capturing the nuanced context in text.
• Efficient Inference: When fine-tuned for tasks like intent classification, these models offer efficient performance suitable for production environments.

BERT vs. DistilBERT:

• BERT: A powerful model with robust performance, but it is computationally intensive and has higher latency, making it less ideal for real-time applications.
• DistilBERT: A distilled version of BERT that retains most of its performance while being significantly smaller and faster.
• Our Choice: DistilBERT was chosen to balance high-quality language understanding with efficiency, making it well-suited for environments where inference speed and reduced resource consumption are critical.

Fine-Tuning Strategy:

• Pre-trained Initialization: Start with a pre-trained model (in this case, DistilBERT) to leverage its extensive language understanding.
• Task-Specific Fine-Tuning: Fine-tune on a labeled dataset that covers the three defined intent categories, allowing the model to adapt its general language knowledge to the specific task.
• Hyperparameter Optimization: Adjust parameters such as learning rate (e.g., 2e-5), batch size (e.g., 32), and the number of epochs (typically around 3) to achieve the best balance between accuracy and computational efficiency.

Model Architecture Details:

• Input Layer: Processes tokenized text using a standard tokenizer compatible with BERT, converting raw text into input IDs and attention masks.
• Transformer Encoder: Uses pre-trained DistilBERT layers to generate contextual embeddings from the input text.
• Classification Head:1) Consists of one or more fully connected (dense) layers that map the encoder’s output to a set of raw scores (logits) corresponding to the three intent classes., 2) Output Logits:These logits are unnormalized scores that indicate how strongly the model associates the input with each class.
• Softmax Activation:Applies the softmax function to the logits to convert them into probability scores that sum to 1, providing an interpretable measure of the model's confidence for each intent class.

5. Dataset and Data Preparation

Dataset Collection:

Historical user queries and manually annotated samples covering three intent categories:

• Consultation & Q&A?(e.g., “What are the mobile network’s strengths?”)
• Ideation & Brainstorming?(e.g., “10 party game ideas”)
• Planning & Scheduling?(e.g., “Plan a 2-day trip to Kyoto”)

Data Labelling:

Each query is tagged with its appropriate intent, as a numeric value, ensuring clear class distinctions.

• 0?for Consultation & Q&A
• 1?for Ideation & Brainstorming
• 2?for Planning & Scheduling

Data Augmentation:

Optionally, generate paraphrases and variations of existing queries to enhance model robustness and improve generalization.

Data Format:

Store data in CSV with at least two columns:

• query:?The raw user text.
• label:?The corresponding numerical label.

Dataset Splits:

1. Training Set: ~80%
2. Validation Set: ~10%
3. Test Set: ~10%

Preprocessing Steps:

• Text Normalization: Convert text to lowercase and remove extraneous whitespace or special characters.
• Tokenization: Use the tokenizer compatible with our pre-trained model (e.g., the DistilBERT tokenizer) to convert text into tokens.
• Padding & Truncation: Ensure all inputs are of a consistent length by padding shorter sequences and truncating longer ones to meet model requirements.

Training Process:

• Environment: Leverage GPU-enabled instances to accelerate the fine-tuning process.
• Loss Function: Utilize cross-entropy loss for multi-class classification, which effectively compares the predicted probabilities against the true labels.
• Validation Strategy: Reserve a hold-out validation set to monitor performance and detect potential overfitting during training.
• Evaluation Metrics: Assess model performance using metrics such as accuracy, precision, recall, and F1-score for a comprehensive evaluation.
• Tools and Frameworks: Implement the training process using frameworks like PyTorch along with Hugging Face’s Transformers library to streamline model integration and fine-tuning.

PyTorch Dataset Class Example

This code creates a dataset class that prepares text data for fine-tuning a DistilBERT model. It tokenizes text inputs, pads or truncates them to a fixed length, and converts both the input data and labels into PyTorch tensors ready for model training.

from transformers import DistilBertTokenizer
from torch.utils.data import Dataset
import torch

class IntentDataset(Dataset):
    def __init__(self, texts, labels, tokenizer: DistilBertTokenizer, max_length: int = 128):
        self.texts = texts
        self.labels = labels
        self.tokenizer = tokenizer
        self.max_length = max_length

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, idx):
        encoding = self.tokenizer(
            self.texts[idx],
            add_special_tokens=True,
            max_length=self.max_length,
            padding='max_length',
            truncation=True,
            return_tensors='pt'
        )
        return {
            'input_ids': encoding['input_ids'].squeeze(),  # Shape: [max_length]
            'attention_mask': encoding['attention_mask'].squeeze(),
            'label': torch.tensor(self.labels[idx], dtype=torch.long)
        }        

6. Model Training and Fine-Tuning

Model Architecture

• Base Model: Start with a pre-trained DistilBERT model (e.g.,?distilbert-base-uncased).
• Classification Head: A fully connected layer added on top of DistilBERT outputs logits for three classes, followed by a softmax activation for probability estimates.

Training Setup

• Framework:?PyTorch
• Loss Function:?Cross-entropy loss for multi-class classification.
• Optimizer:?AdamW optimizer.
• Hyperparameters:Learning Rate:?2e-5Batch Size:?32Epochs:?Typically 3 (subject to tuning based on validation performance)

Training Script Example

import torch
from transformers import DistilBertForSequenceClassification, AdamW
from torch.utils.data import DataLoader

# Initialize tokenizer and model
tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')
model = DistilBertForSequenceClassification.from_pretrained('distilbert-base-uncased', num_labels=3)

# Assume train_texts and train_labels are loaded from CSV
train_dataset = IntentDataset(train_texts, train_labels, tokenizer)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)

# Set up optimizer
optimizer = AdamW(model.parameters(), lr=2e-5)
model.train()

num_epochs = 3
for epoch in range(num_epochs):
    epoch_loss = 0
    for batch in train_loader:
        optimizer.zero_grad()
        outputs = model(
            input_ids=batch['input_ids'],
            attention_mask=batch['attention_mask'],
            labels=batch['label']
        )
        loss = outputs.loss
        loss.backward()
        optimizer.step()
        epoch_loss += loss.item()
    
    avg_loss = epoch_loss / len(train_loader)
    print(f"Epoch {epoch + 1}/{num_epochs}, Loss: {avg_loss:.4f}")

# Save the trained model
torch.save(model.state_dict(), "intent_extraction_distilbert.pt")        

7. Inference Service with FastAPI

REST API Design

• Endpoint:?POST /api/v1/intent
• Request Payload Example:

{
    "query": "Plan a 2-day trip to Kyoto"
}        

• Response Payload Example:

{
    "intent": "Planning & Scheduling",
    "confidence": 0.92
}        

FastAPI Implementation Example

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import torch
from transformers import DistilBertTokenizer, DistilBertForSequenceClassification

app = FastAPI()

# Load tokenizer and fine-tuned DistilBERT model on startup
tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')
model = DistilBertForSequenceClassification.from_pretrained('path/to/fine_tuned_model')
model.load_state_dict(torch.load("intent_extraction_distilbert.pt", map_location=torch.device('cpu')))
model.eval()  # Set model to evaluation mode

# Mapping from numerical labels to intents
intent_labels = {
    0: "Consultation & Q&A",
    1: "Ideation & Brainstorming",
    2: "Planning & Scheduling"
}

class QueryRequest(BaseModel):
    query: str

@app.post("/api/v1/intent")
async def extract_intent(request: QueryRequest):
    try:
        # Preprocess the query
        inputs = tokenizer(
            request.query,
            add_special_tokens=True,
            max_length=128,
            padding='max_length',
            truncation=True,
            return_tensors='pt'
        )
        # Inference
        with torch.no_grad():
            outputs = model(**inputs)
        logits = outputs.logits
        probabilities = torch.softmax(logits, dim=1)
        confidence, pred_class = torch.max(probabilities, dim=1)
        intent = intent_labels[pred_class.item()]
        return {"intent": intent, "confidence": confidence.item()}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
        

8. Containerization and AWS Deployment

Dockerfile Example

FROM python:3.9-slim

# Set working directory
WORKDIR /app

# Install dependencies
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# Copy source code
COPY . .

# Expose port for FastAPI
EXPOSE 8000

# Start the FastAPI service using Uvicorn
CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"]        

requirements.txt Example

fastapi
uvicorn
torch
transformers
pydantic        

AWS Cloud Deployment

Container Registry and Orchestration

• ECR:?Build, tag, and push the Docker image to AWS Elastic Container Registry.
• EKS/ECS:?Deploy the containerized service using AWS EKS (Kubernetes) or ECS with Fargate.

Deployment Steps

1. Build and Tag Docker Image:

docker build -t intent-extraction-service .
docker tag intent-extraction-service:latest <aws_account_id>.dkr.ecr.<region>.amazonaws.com/intent-extraction-service:latest        

2. Push to ECR:

docker push <aws_account_id>.dkr.ecr.<region>.amazonaws.com/intent-extraction-service:latest        

3. Deploy on AWS:

• EKS: Use a Kubernetes Deployment manifest to specify the container image, resource limits, and environment variables.
• ECS: Create a task definition and service with auto-scaling policies.
• Networking: Utilize an Application Load Balancer (ALB) for traffic routing.

Example Kubernetes Deployment Manifest

apiVersion: apps/v1
kind: Deployment
metadata:
  name: intent-extraction-deployment
spec:
  replicas: 3
  selector:
    matchLabels:
      app: intent-extraction
  template:
    metadata:
      labels:
        app: intent-extraction
    spec:
      containers:
      - name: intent-extraction-container
        image: <aws_account_id>.dkr.ecr.<region>.amazonaws.com/intent-extraction-service:latest
        ports:
        - containerPort: 8000
        resources:
          requests:
            cpu: "250m"
            memory: "512Mi"
          limits:
            cpu: "500m"
            memory: "1Gi"        

CI/CD Pipeline

• Source Control:?Git-based repository.
• Automation:?Use AWS CodePipeline and CodeBuild for automated testing, image building, and deployment.
• Testing:?Unit tests for the API endpoints and integration tests for the model inference pipeline.

Logging, Monitoring, and Security

Logging and Monitoring

• Logging: Use structured JSON logging in FastAPI. Forward logs to AWS CloudWatch Logs.
• Monitoring: Track metrics like response times, throughput, and error rates with CloudWatch.
• Alerting: Set up alarms for performance degradation or spikes in error rates.

Security Considerations

• Communication: All API communications use HTTPS.
• IAM Policies: Apply least privilege principles for AWS resource access.
• Secrets Management: Manage sensitive data using AWS Secrets Manager.

9. Future Enhancements

• Model Updates: Implementing a pipeline for periodic retraining means setting up an automated or semi-automated system that regularly re-trains the model as new query data becomes available. This ensures that the model stays current with evolving language patterns and user behaviors, improving its accuracy and relevance over time.
• Contextual Analysis: This enhancement involves moving beyond analyzing single, isolated queries. By integrating multi-turn conversation analysis—where the system takes into account previous interactions or session context—the model can better understand ambiguous queries. Incorporating the context from earlier conversation turns can lead to a more nuanced interpretation of the user's intent.
• Fallback Mechanism: A fallback mechanism is designed to handle instances when the model’s confidence in its prediction is low. Instead of forcing a potentially incorrect classification, the system can prompt the user for clarification or take alternative actions (such as routing the query to a human agent). This approach helps maintain a good user experience even when the model is uncertain.
• A/B Testing: A/B testing involves deploying a new version of the model or changes to intent definitions alongside the current system. By comparing performance metrics (like accuracy, user engagement, or conversion rates) between the two versions, you can objectively assess whether the new improvements provide a significant benefit before fully rolling them out.
• Metadata Tagging: We can extend the service by adding a metadata tagging branch to the DistilBERT model so that it not only classifies the overall intent but also extracts metadata like planning type (e.g., travel or events) and implied constraints (duration, location, preferences). This extension involves augmenting the training dataset with metadata labels, modifying the model architecture for multi-task learning with parallel output heads, and updating the API response to include both intent and metadata information.

10. Conclusion

This design outlines a robust, scalable, and cost-efficient intent extraction service for a conversational AI feature. By leveraging a fine-tuned BERT-based model, the system can accurately classify user queries across consultation, ideation, and planning use cases. The deployment strategy using RESTful APIs within a Kubernetes environment ensures that the service remains responsive and scalable, meeting both current and future demands.