Building a minimal production-ready sentiment analysis model
What is a production-ready model?
We hear a lot about productionized machine learning, but what does it really mean to have a model that can thrive in real-world applications?There are plenty of things that go into, and contribute, to the efficacy of a machine learning model in production. For the sake of this article we will be focusing on five of them.
- Reproducibility
- Monitoring
- Testing
- Automation
- Version Control
Serving Inferences
The most important part of building a production-ready machine learning model is being able to access it.
For this purpose, we build a fastapi client that serves sentiment analysis responses. We utilize pydantic to ensure structure for the input and output. The model that we use is the base sentiment analysis pipeline from huggingface’s transformers library, allowing us to begin testing with a pre-trained model.
# Filename: main.py
from fastapi import FastAPI
from pydantic import BaseModel
from transformers import pipeline
app = FastAPI()
classifier = pipeline("sentiment-analysis")
class TextInput(BaseModel):
text: str
class SentimentOutput(BaseModel):
text: str
sentiment: str
score: float
@app.post("/predict", response_model=SentimentOutput)
async def predict_sentiment(input_data: TextInput):
result = classifier(input_data.text)[0]
return SentimentOutput(
text=input_data.text,
sentiment=result["label"],
score=result["score"]
)
To ensure that our work is reproducible, we can use a requirements.txt file and pip.
# Filename: requirements.txt
# Note: This has all required packages for the final result.
fastapi==0.68.1
uvicorn==0.15.0
transformers==4.30.0
torch==2.0.0
pydantic==1.10.0
numpy==1.24.3
sentencepiece==0.1.99
protobuf==3.20.3
prometheus-client==0.17.1
To install this, initialize venv in your files and run:pip install -r requirements.txt.
To host this API simply run: uvicorn main:app –reload.
Now you have an api that you can query using:
curl -X POST "http://localhost:8000/predict" \
-H "Content-Type: application/json" \
-d '{"text": "I love using FastAPI!"}'
or any API tool you wish (i.e. Postman). You should get a result back that includes the text query, the sentiment predicted, and the confidence of the prediction.
We will be using GitHub for CI/CD later, so I would recommend initializing and using git in this directory.
We now have a locally hosted machine learning inference API.
Further Improving Reproducibility
To allow our code to have more consistent execution, we will utilize Docker. Docker simulates a lightweight environment that allows applications to run in isolated containers, similar to virtual machines. This isolation ensures that applications can execute consistently across any computer with Docker installed, regardless of the underlying system.
Firstly, set up Docker for your given operating system.
# Filename: Dockerfile
# Use the official Python 3.9 slim image as the base
FROM python:3.9-slim
# Set the working directory inside the container to /app
WORKDIR /app
# Copy the requirements.txt file to the working directory
COPY requirements.txt .
# Install the Python dependencies listed in requirements.txt
RUN pip install -r requirements.txt
# Copy the main application file (main.py) to the working directory
COPY main.py .
# Define the command to run the FastAPI application with Uvicorn
CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"]
At this point, you should have the directory as below.
your-project/
├── Dockerfile
├── requirements.txt
└── main.py
Now, you can build the image and run this API using:
# Build the Docker image
docker build -t sentiment-api .
# Run the container
docker run -p 8000:8000 sentiment-api
You should now be able to query just as you did before.
curl -X POST "http://localhost:8000/predict" \
-H "Content-Type: application/json" \
-d '{"text": "I love using FastAPI!"}'
We now have a containerized, locally hosted machine learning inference API.
Adding Basic Monitoring
In machine learning applications, monitoring is crucial for understanding model performance and ensuring it meets expected accuracy and efficiency. Tools like Prometheus help track metrics such as prediction latency, request counts, and model output distributions, enabling you to identify issues like model drift or resource bottlenecks. This proactive approach ensures that your ML models remain effective over time and can adapt to changing data or usage patterns. In our case, we are focused on prediction time, requests, and gathering information about our queries.
from fastapi import FastAPI
from pydantic import BaseModel
from transformers import pipeline
from prometheus_client import Counter, Histogram, start_http_server
import time
# Start prometheus metrics server on port 8001
start_http_server(8001)
app = FastAPI()
# Metrics
PREDICTION_TIME = Histogram('prediction_duration_seconds', 'Time spent processing prediction')
REQUESTS = Counter('prediction_requests_total', 'Total requests')
SENTIMENT_SCORE = Histogram('sentiment_score', 'Histogram of sentiment scores', buckets=[0.0, 0.25, 0.5, 0.75, 1.0])
class TextInput(BaseModel):
text: str
class SentimentOutput(BaseModel):
text: str
sentiment: str
score: float
@app.post("/predict", response_model=SentimentOutput)
async def predict_sentiment(input_data: TextInput):
REQUESTS.inc()
start_time = time.time()
result = classifier(input_data.text)[0]
score = result["score"]
SENTIMENT_SCORE.observe(score) # Record the sentiment score
PREDICTION_TIME.observe(time.time() - start_time)
return SentimentOutput(
text=input_data.text,
sentiment=result["label"],
score=score
)
Utilizing a Custom Model
While the process of building and fine-tuning a model is not the intent of this project, it is important to understand how a model can be added to this process.
# Filename: train.py
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from datasets import load_dataset
from torch.utils.data import DataLoader
def train_model():
# Load dataset
full_dataset = load_dataset("stanfordnlp/imdb", split="train")
dataset = full_dataset.shuffle(seed=42).select(range(10000))
model_name = "distilbert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2)
optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5)
# Use GPU if available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
model.train()
# Create a DataLoader for batching
dataloader = DataLoader(dataset, batch_size=8, shuffle=True)
# Training loop
num_epochs = 3 # Set the number of epochs
for epoch in range(num_epochs):
total_loss = 0
for batch in dataloader:
inputs = tokenizer(batch["text"], truncation=True, padding=True, return_tensors="pt", max_length=512).to(device)
labels = torch.tensor(batch["label"]).to(device)
optimizer.zero_grad()
outputs = model(**inputs, labels=labels)
loss = outputs.loss
loss.backward()
optimizer.step()
total_loss += loss.item()
avg_loss = total_loss / len(dataloader)
print(f"Epoch {epoch + 1}/{num_epochs}, Loss: {avg_loss:.4f}")
# Save the model
model.save_pretrained("./model/")
tokenizer.save_pretrained("./model/")
# Test the model with sample sentences
test_sentences = [
"This movie was fantastic!",
"I absolutely hated this film.",
"It was just okay, not great.",
"An absolute masterpiece!",
"Waste of time!",
"A beautiful story and well acted.",
"Not my type of movie.",
"It could have been better.",
"A thrilling adventure from start to finish!",
"Very disappointing."
]
# Switch model to evaluation mode
model.eval()
# Prepare tokenizer for test inputs
inputs = tokenizer(test_sentences, truncation=True, padding=True, return_tensors="pt", max_length=512).to(device)
with torch.no_grad():
outputs = model(**inputs)
predictions = torch.argmax(outputs.logits, dim=1)
# Print predictions
for sentence, prediction in zip(test_sentences, predictions):
sentiment = "positive" if prediction.item() == 1 else "negative"
print(f"Input: \"{sentence}\" -> Predicted sentiment: {sentiment}")
# Call the function to train the model and test it
train_model()
To make sure that we can query our new model that we have trained we have to update a few of our existing files. For instance, in main.py we now use the model from ./model and load it as a pretrained model. Additionally, for comparison’s sake, we add now have two endpoints to use, /predict/naive and predict/trained.
# Filename: main.py
from fastapi import FastAPI
from pydantic import BaseModel
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from transformers import pipeline
from prometheus_client import Counter, Histogram, start_http_server
import time
# Start prometheus metrics server on port 8001
start_http_server(8001)
app = FastAPI()
# Load the trained model and tokenizer from the local directory
model_path = "./model" # Path to your saved model
tokenizer = AutoTokenizer.from_pretrained(model_path)
trained_model = AutoModelForSequenceClassification.from_pretrained(model_path)
# Create pipelines
naive_classifier = pipeline("sentiment-analysis", device=-1)
trained_classifier = pipeline("sentiment-analysis", model=trained_model, tokenizer=tokenizer, device=-1)
# Metrics
PREDICTION_TIME = Histogram('prediction_duration_seconds', 'Time spent processing prediction')
REQUESTS = Counter('prediction_requests_total', 'Total requests')
SENTIMENT_SCORE = Histogram('sentiment_score', 'Histogram of sentiment scores', buckets=[0.0, 0.25, 0.5, 0.75, 1.0])
class TextInput(BaseModel):
text: str
class SentimentOutput(BaseModel):
text: str
sentiment: str
score: float
@app.post("/predict/naive", response_model=SentimentOutput)
async def predict_naive_sentiment(input_data: TextInput):
REQUESTS.inc()
start_time = time.time()
result = naive_classifier(input_data.text)[0]
score = result["score"]
SENTIMENT_SCORE.observe(score) # Record the sentiment score
PREDICTION_TIME.observe(time.time() - start_time)
return SentimentOutput(
text=input_data.text,
sentiment=result["label"],
score=score
)
@app.post("/predict/trained", response_model=SentimentOutput)
async def predict_trained_sentiment(input_data: TextInput):
REQUESTS.inc()
start_time = time.time()
result = trained_classifier(input_data.text)[0]
score = result["score"]
SENTIMENT_SCORE.observe(score) # Record the sentiment score
We also must update our Dockerfile to include our model files.
# Filename: Dockerfile
FROM python:3.9-slim
WORKDIR /app
COPY requirements.txt .
RUN pip install -r requirements.txt
COPY main.py .
COPY ./model ./model
CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"]
Importantly, if you are using git, make sure that you add the pytorch_model.bin file to git lfs, so that you can push to GitHub. git lfs allows you to use version control on very large files.
Adding Testing and CI/CD
CI/CD and testing streamline the deployment of machine learning models by ensuring that code changes are automatically integrated, tested, and deployed, which reduces the risk of errors and enhances model reliability. This process promotes continuous improvement and faster iteration cycles, allowing teams to deliver high-quality, production-ready models more efficiently. Firstly, we create two very basic tests to ensure that our model is performing acceptably.
# Filename: test_model.py
import pytest
from fastapi.testclient import TestClient
from main import app
client = TestClient(app)
def test_positive_sentiment():
response = client.post(
"/predict/trained",
json={"text": "This is amazing!"}
)
assert response.status_code == 200
data = response.json()
assert data["sentiment"] == "LABEL_1"
assert data["score"] > 0.5
def test_negative_sentiment():
response = client.post(
"/predict/trained",
json={"text": "This is terrible!"}
)
assert response.status_code == 200
data = response.json()
assert data["sentiment"] == "LABEL_0"
assert data["score"] < 0.5
To test your code, you can simply run pytest or python -m pytest while your endpoint is running.
However, we will add automated testing CI/CD (continuous integration and continuous delivery) when pushed to GitHub.
# Filename: .github/workflows/ci_cd.yml
name: CI/CD
on: [push]
jobs:
test:
runs-on: ubuntu-latest
steps:
- name: Checkout code
uses: actions/checkout@v2
with:
lfs: true
- name: Set up Python
uses: actions/setup-python@v2
with:
python-version: '3.9'
- name: Install dependencies
run: |
pip install -r requirements.txt
pip install pytest httpx
- name: Run tests
run: pytest
Our final project structure should appear as below.
sentiment-analysis-project/
├── .github/
│ └── workflows/
│ └── ci_cd.yml
├── test_model.py
├── main.py
├── Dockerfile
├── requirements.txt
└── train.py
Now, whenever we push to GitHub, it will run an automated process that checks out the code, sets up a Python 3.9 environment, installs dependencies, and runs our tests using pytest.
Conclusion
In this project, we’ve developed a production-ready sentiment analysis API that highlights key aspects of deploying machine learning models. While it doesn’t encompass every facet of the field, it provides a representative sampling of essential tasks involved in the process. By examining these components, I hope to clarify concepts you may have encountered but weren’t quite sure how they fit together in a practical setting.
Minimum Viable MLE was originally published in Towards Data Science on Medium, where people are continuing the conversation by highlighting and responding to this story.