What You'll Build

A production-grade sentiment classifier trained on the IMDB movie review dataset (50,000 reviews) that:

Dataset: IMDB Large Movie Review Dataset — load directly from HuggingFace: datasets.load_dataset("imdb"). No manual download needed.

Project Setup

python -m venv .venv
source .venv/bin/activate

pip install pandas numpy scikit-learn matplotlib seaborn \
            transformers datasets torch accelerate flask

Create sentiment.ipynb for exploration and app.py for the Flask server.

Step 1 — Load the Dataset

from datasets import load_dataset

ds = load_dataset("imdb")
print(ds)
# DatasetDict with train (25k) and test (25k) splits

# Convert to pandas for easier manipulation
import pandas as pd
train_df = pd.DataFrame(ds['train'])
test_df  = pd.DataFrame(ds['test'])

print(train_df.head(3))
print(train_df['label'].value_counts())  # balanced: 12,500 pos / 12,500 neg

Step 2 — Text Cleaning

import re

def clean_text(text):
    text = re.sub(r'<.*?>', ' ', text)         # strip HTML tags
    text = re.sub(r'[^a-zA-Z\s]', ' ', text)  # keep only letters
    text = text.lower().strip()
    text = re.sub(r'\s+', ' ', text)           # collapse whitespace
    return text

train_df['text_clean'] = train_df['text'].apply(clean_text)
test_df['text_clean']  = test_df['text'].apply(clean_text)

Step 3 — TF-IDF Baseline

from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import classification_report, ConfusionMatrixDisplay

vectorizer = TfidfVectorizer(max_features=30_000, ngram_range=(1, 2),
                              min_df=2, sublinear_tf=True)

X_train = vectorizer.fit_transform(train_df['text_clean'])
X_test  = vectorizer.transform(test_df['text_clean'])
y_train = train_df['label']
y_test  = test_df['label']

lr = LogisticRegression(max_iter=1000, C=1.0)
lr.fit(X_train, y_train)
y_pred = lr.predict(X_test)

print(classification_report(y_test, y_pred, target_names=['Negative','Positive']))

Expected output (approximately):

              precision    recall  f1-score   support
    Negative       0.91      0.93      0.92     12500
    Positive       0.93      0.91      0.92     12500
    accuracy                           0.92     25000
Why bigrams? ngram_range=(1,2) captures two-word phrases like "not good" or "highly recommend" — these flip sentiment and a unigram model misses them entirely.

Confusion Matrix

import matplotlib.pyplot as plt

ConfusionMatrixDisplay.from_predictions(
    y_test, y_pred,
    display_labels=['Negative', 'Positive'],
    cmap='Blues'
)
plt.title('Logistic Regression Confusion Matrix')
plt.show()

Step 4 — BERT Fine-Tuning

from transformers import (AutoTokenizer, AutoModelForSequenceClassification,
                          TrainingArguments, Trainer)
from datasets import Dataset
import numpy as np

MODEL = "distilbert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(MODEL)

def tokenize(batch):
    return tokenizer(batch['text'], truncation=True, padding='max_length',
                     max_length=256)

# Use a 5k subset for fast training — increase for full accuracy
train_small = ds['train'].shuffle(seed=42).select(range(5000))
test_small  = ds['test'].shuffle(seed=42).select(range(1000))

train_tok = train_small.map(tokenize, batched=True)
test_tok  = test_small.map(tokenize, batched=True)
train_tok = train_tok.rename_column("label", "labels")
test_tok  = test_tok.rename_column("label", "labels")
train_tok.set_format("torch", columns=["input_ids","attention_mask","labels"])
test_tok.set_format("torch", columns=["input_ids","attention_mask","labels"])
model = AutoModelForSequenceClassification.from_pretrained(MODEL, num_labels=2)

def compute_metrics(eval_pred):
    logits, labels = eval_pred
    preds = np.argmax(logits, axis=-1)
    acc = (preds == labels).mean()
    return {"accuracy": acc}

args = TrainingArguments(
    output_dir="./bert-sentiment",
    num_train_epochs=3,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=32,
    evaluation_strategy="epoch",
    save_strategy="no",
    load_best_model_at_end=False,
    logging_steps=50,
    report_to="none",
)

trainer = Trainer(
    model=model,
    args=args,
    train_dataset=train_tok,
    eval_dataset=test_tok,
    compute_metrics=compute_metrics,
)

trainer.train()
results = trainer.evaluate()
print(f"BERT accuracy: {results['eval_accuracy']:.4f}")  # ~0.94–0.95
GPU recommended: Training on 5k examples takes ~5 minutes on a GPU (Colab T4 works). On CPU it will take 30–60 minutes. Use model.to("cuda") if you have a local GPU.

Step 5 — Model Comparison

ModelAccuracyTraining TimeInference Speed
TF-IDF + Logistic Regression~92%30 secFast
DistilBERT fine-tuned~95%~5 min (GPU)Slower

For most production use cases, the TF-IDF baseline is often "good enough" and much faster to serve. Use BERT when accuracy is critical or when the text is more complex than movie reviews.

Step 6 — Flask API

Create app.py:

from flask import Flask, request, jsonify
import pickle, re
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression

app = Flask(__name__)

# Load pre-trained objects (save them first with pickle)
# with open('tfidf.pkl','wb') as f: pickle.dump(vectorizer, f)
# with open('lr_model.pkl','wb') as f: pickle.dump(lr, f)
with open('tfidf.pkl','rb')    as f: vec = pickle.load(f)
with open('lr_model.pkl','rb') as f: clf = pickle.load(f)

def clean(text):
    text = re.sub(r'<.*?>', ' ', text)
    text = re.sub(r'[^a-zA-Z\s]', ' ', text)
    return text.lower().strip()

@app.route('/predict', methods=['POST'])
def predict():
    data = request.get_json()
    text = data.get('text', '')
    if not text:
        return jsonify({'error': 'No text provided'}), 400
    features = vec.transform([clean(text)])
    label = int(clf.predict(features)[0])
    prob  = clf.predict_proba(features)[0][label]
    return jsonify({
        'label':      'positive' if label == 1 else 'negative',
        'confidence': round(float(prob), 4)
    })

if __name__ == '__main__':
    app.run(debug=True, port=5000)
# Test the API
curl -X POST http://localhost:5000/predict \
  -H "Content-Type: application/json" \
  -d '{"text": "This movie was absolutely fantastic — loved every minute."}'
# {"confidence": 0.9987, "label": "positive"}

What to Try Next