This page builds a binary sentiment classifier by fine-tuning a pre-trained BERT encoder (bert-base-uncased) with a small two-layer head in PyTorch, saves the weights, and then reloads them to score new sentences. It is a compact, end-to-end template for any text classification task: swap the dataset and the label count and the same structure applies.
Sentiment analysis maps a piece of text to a polarity label. Rather than training word embeddings from scratch, the model reuses BERT, a transformer encoder pre-trained on large unlabeled corpora with masked language modelling. BERT turns each token into a contextual 768-dimensional vector; the pooler_output is a learned projection of the special [CLS] token and acts as a sentence-level summary. A small feed-forward head (fc1 to 128 units with ReLU, then fc2 to a single logit) learns to separate positive from negative on top of that summary. Because all BERT parameters are passed to the optimizer, the whole network is fine-tuned, not just the head.
The tokenizer splits text into WordPiece sub-words, adds [CLS] and [SEP], pads or truncates to 128 tokens and returns an attention_mask so padding positions are ignored by self-attention. The encoder produces contextual vectors, the pooler summarises them, and the head reduces the summary to one number. A single logit with a sigmoid is the standard formulation for two classes; for three or more classes you would output num_labels logits and use CrossEntropyLoss instead.
#!/usr/bin/env python
import torch
from torch.utils.data import Dataset, DataLoader
from torch import nn
from transformers import BertTokenizer, BertModel
from sklearn.metrics import accuracy_score, classification_report
# Define the TextDataset class
class TextDataset(Dataset):
def __init__(self, texts, labels):
self.texts = texts
self.labels = labels
self.tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
def __len__(self):
return len(self.texts)
def __getitem__(self, idx):
text = self.texts[idx]
label = self.labels[idx]
tokens = self.tokenizer(
text,
padding='max_length',
max_length=128,
truncation=True,
return_tensors='pt'
)
input_ids = tokens['input_ids'].squeeze(0)
attention_mask = tokens['attention_mask'].squeeze(0)
return {
'input_ids': input_ids,
'attention_mask': attention_mask,
'labels': torch.tensor(label, dtype=torch.float)
}
# Expanded dataset with both positive and negative examples
text_data = [
"This movie is amazing!", # positive
"I really disliked the plot.", # negative
"The acting was superb.", # positive
"This book is a masterpiece.", # positive
"What a terrible waste of time.", # negative
"I love this!", # positive
"This was the worst movie ever.", # negative
"Fantastic experience, loved it.", # positive
"Not good at all.", # negative
"Absolutely fantastic.", # positive
"I hated every minute of it.", # negative
"It was a decent watch.", # neutral/positive
"Great movie!", # positive
"I will never watch this again.", # negative
"Loved every second of it.", # positive
"Awful plot, horrible acting.", # negative
"Brilliant direction and storytelling.",# positive
"Worst experience I've had.", # negative
"The movie was quite boring.", # negative
"Best performance of the year." # positive
]
labels = [1, 0, 1, 1, 0, 1, 0, 1, 0, 1, 0, 1, 1, 0, 1, 0, 1, 0, 0, 1]
# Split into train and test sets (80% training, 20% testing)
train_size = int(0.8 * len(text_data))
train_dataset = TextDataset(text_data[:train_size], labels[:train_size])
test_dataset = TextDataset(text_data[train_size:], labels[train_size:])
train_dataloader = DataLoader(train_dataset, batch_size=4)
test_dataloader = DataLoader(test_dataset, batch_size=4)
# Define the Model with two fully connected layers
class SentimentClassifier(nn.Module):
def __init__(self):
super().__init__()
self.bert = BertModel.from_pretrained('bert-base-uncased')
self.fc1 = nn.Linear(self.bert.config.hidden_size, 128)
self.fc2 = nn.Linear(128, 1)
def forward(self, input_ids, attention_mask):
outputs = self.bert(input_ids, attention_mask)
pooled_output = outputs.pooler_output
x = torch.relu(self.fc1(pooled_output)) # Add non-linearity
output = self.fc2(x)
return output # Return raw logits, no sigmoid here
# Initialize the Model
model = SentimentClassifier()
# Loss Function and Optimizer
criterion = nn.BCEWithLogitsLoss() # Use BCEWithLogitsLoss
optimizer = torch.optim.Adam(model.parameters(), lr=3e-5)
# Train the model or load the existing model
train_model = True # Set to False if you want to skip training and load the model
if train_model:
# Training Loop
num_epochs = 3
for epoch in range(num_epochs):
model.train() # Ensure the model is in training mode
for batch in train_dataloader:
input_ids = batch['input_ids']
attention_mask = batch['attention_mask']
labels = batch['labels'].float()
optimizer.zero_grad()
# Get model outputs
outputs = model(input_ids, attention_mask)
# Squeeze the outputs to remove the extra dimension
outputs = outputs.squeeze(1)
# Compute the loss
loss = criterion(outputs, labels)
# Backpropagation
loss.backward()
optimizer.step()
print(f'Epoch [{epoch+1}/{num_epochs}]')
print(f'Loss: {loss.item()}')
# Save the trained model
torch.save(model.state_dict(), 'sentiment_model.pth')
print("Model saved as 'sentiment_model.pth'")
else:
# Load the saved model
model.load_state_dict(torch.load('sentiment_model.pth'))
model.eval() # Set to evaluation mode
print("Model loaded from 'sentiment_model.pth'")
# Evaluation
model.eval() # Set the model to evaluation mode
all_preds = []
all_labels = []
threshold = 0.5 # Set a threshold for classification
with torch.no_grad():
for batch in test_dataloader:
input_ids = batch['input_ids']
attention_mask = batch['attention_mask']
outputs = model(input_ids, attention_mask)
# Squeeze the outputs and apply sigmoid + threshold to get predictions
outputs = torch.sigmoid(outputs.squeeze(1))
preds = (outputs > threshold).float()
all_preds.extend(preds.cpu().numpy())
all_labels.extend(batch['labels'].cpu().numpy())
# Print accuracy and classification report
print(f'Accuracy: {accuracy_score(all_labels, all_preds)}')
print(classification_report(all_labels, all_preds))
# Optional: Print raw outputs for debugging
with torch.no_grad():
for batch in test_dataloader:
input_ids = batch['input_ids']
attention_mask = batch['attention_mask']
outputs = model(input_ids, attention_mask)
# Print raw outputs to inspect the values
print("Model outputs:", outputs)
preds = (torch.sigmoid(outputs.squeeze(1)) > threshold).float()
print("Predictions:", preds)
#!/usr/bin/env python
import torch
from torch import nn
from transformers import BertTokenizer, BertModel
from prettytable import PrettyTable
# Define the Model with two fully connected layers
class SentimentClassifier(nn.Module):
def __init__(self):
super().__init__()
self.bert = BertModel.from_pretrained('bert-base-uncased')
self.fc1 = nn.Linear(self.bert.config.hidden_size, 128)
self.fc2 = nn.Linear(128, 1)
def forward(self, input_ids, attention_mask):
outputs = self.bert(input_ids, attention_mask)
pooled_output = outputs.pooler_output
x = torch.relu(self.fc1(pooled_output)) # Add non-linearity
output = self.fc2(x)
return output # Return raw logits
# Load the trained model
model = SentimentClassifier()
model.load_state_dict(torch.load('sentiment_model.pth'))
model.eval() # Set the model to evaluation mode
# Load the tokenizer
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
# Example positive and negative statements
statements = [
"I absolutely love this movie!", # Positive
"This is the worst film I have ever seen.", # Negative
"The book was fantastic!", # Positive
"I really hated the acting.", # Negative
"Amazing story and great performances.", # Positive
"It was a complete waste of time.", # Negative
]
# Prepare a table using PrettyTable
table = PrettyTable()
table.field_names = ["Statement", "Prediction"]
# Run inference on each statement
threshold = 0.5 # Classification threshold
for statement in statements:
# Tokenize and preprocess the input statement
tokens = tokenizer(
statement,
padding = 'max_length',
max_length = 128,
truncation = True,
return_tensors = 'pt'
)
input_ids = tokens['input_ids']
attention_mask = tokens['attention_mask']
# Run the model and get predictions
with torch.no_grad():
outputs = model(input_ids, attention_mask)
outputs = torch.sigmoid(outputs).squeeze(1) # Apply sigmoid to logits
prediction = (outputs > threshold).float().item() # Apply threshold
# Convert the prediction to "Positive" or "Negative"
sentiment = "Positive" if prediction == 1.0 else "Negative"
# Add the statement and its prediction to the table
table.add_row([statement, sentiment])
# Print the table
print(table)
+-----------------------------------------------+------------+
| Statement | Prediction |
+-----------------------------------------------+------------+
| I absolutely love this movie! | Positive |
| This is the worst film I have ever seen. | Negative |
| The book was fantastic! | Positive |
| I really hated the acting. | Negative |
| Amazing story and great performances. | Positive |
| It was a complete waste of time. | Negative |
+-----------------------------------------------+------------+
Each iteration pulls a batch from the DataLoader, clears gradients from the previous step with optimizer.zero_grad(), runs the forward pass, computes binary cross-entropy on the logits, back-propagates and updates every parameter with Adam. The learning rate of 3e-5 is in the range commonly used for BERT fine-tuning: large rates destroy the pre-trained weights, a phenomenon usually called catastrophic forgetting. model.train() enables dropout inside BERT; model.eval() together with torch.no_grad() disables dropout and gradient tracking for evaluation, which is both faster and deterministic.
The script prints only the loss of the last batch of each epoch, which is noisy. In practice you would accumulate the mean loss over the epoch and also track validation loss, so you can stop training when validation stops improving.
python -m venv .venv && . .venv/bin/activate
pip install torch transformers scikit-learn prettytable
python build_sentiment_model.py # trains 3 epochs, writes sentiment_model.pth
python use_sentiment_model.py # reloads weights and scores new statementsThe inference script must define the same SentimentClassifier class, because state_dict() stores only tensors keyed by parameter name, not code. Loading on a CPU-only machine weights saved on a GPU needs torch.load(path, map_location="cpu"). A minimal scoring function looks like this:
def predict(text, model, tokenizer, threshold=0.5):
enc = tokenizer(text, padding="max_length", max_length=128,
truncation=True, return_tensors="pt")
with torch.no_grad():
logit = model(enc["input_ids"], enc["attention_mask"]).squeeze()
p = torch.sigmoid(logit).item()
return ("positive" if p > threshold else "negative"), round(p, 3)| Approach | Strengths | Weaknesses |
|---|---|---|
| TF-IDF + logistic regression | Fast, tiny, explainable, trains in seconds | Misses negation and context ("not good") |
| Fine-tuned BERT (this page) | Strong accuracy with modest labelled data; understands context | ~110M parameters, needs a GPU for comfortable training, slower inference |
| DistilBERT / smaller encoders | Roughly half the size, faster inference | Small accuracy drop versus full BERT |
| Prompting a large LLM | No training data needed | Per-call cost and latency, less consistent labels |
Choose fine-tuned BERT when you have at least a few thousand labelled examples, need consistent labels at scale, and can afford a GPU for training. For quick baselines or very constrained hardware, a linear model on TF-IDF features is often good enough and much easier to operate.
DataLoader(train_dataset, batch_size=4) does not shuffle; pass shuffle=True for training so batches are not ordered by how the list was written.TextDataset.__init__ is fine once, but avoid creating datasets repeatedly in a loop.forward and using BCEWithLogitsLoss is correct; adding a sigmoid in the model as well would silently squash gradients..to(device).