machine learning
deep learning
document classification
supervised learning
neural networks

Supervised learningdocument classification using deep learning techniques

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

Document classification is the task of assigning a label to a piece of text, such as spam versus not spam, support ticket category, or news topic. In supervised learning, the model learns that mapping from examples where each document already has the correct label.

What Deep Learning Changes

Older text-classification pipelines often relied on manual feature engineering, such as bag-of-words counts or hand-built keyword rules. Deep learning replaces much of that manual work by learning vector representations and decision boundaries directly from the training data.

In practice, the model pipeline usually has four stages:

  • Convert raw text into tokens or subword units
  • Turn tokens into numeric vectors
  • Aggregate contextual information across the document
  • Predict one label or a probability distribution across labels

Common model families include:

  • CNNs for local phrase patterns
  • LSTMs or GRUs for sequence-aware modeling
  • Transformer encoders for richer contextual representations

The right choice depends on data size, latency requirements, and how much training compute you can afford.

A Small Runnable Example

The example below uses TensorFlow and Keras to build a tiny text classifier. It is intentionally small so the code is easy to run and understand, but the same pattern scales to larger datasets.

python
1import tensorflow as tf
2from tensorflow.keras import layers
3
4texts = [
5    "team won the championship after a close match",
6    "player scored three goals in the final game",
7    "government announced a new tax policy today",
8    "parliament debated the budget proposal all evening",
9    "new graphics card improves gaming performance",
10    "software update fixes security bugs in the browser",
11]
12
13labels = [0, 0, 1, 1, 2, 2]
14label_names = ["sports", "politics", "technology"]
15
16dataset = tf.data.Dataset.from_tensor_slices((texts, labels)).batch(2)
17
18vectorizer = layers.TextVectorization(
19    max_tokens=5000,
20    output_mode="int",
21    output_sequence_length=20,
22)
23vectorizer.adapt(tf.data.Dataset.from_tensor_slices(texts).batch(2))
24
25model = tf.keras.Sequential([
26    vectorizer,
27    layers.Embedding(input_dim=5000, output_dim=32, mask_zero=True),
28    layers.Bidirectional(layers.LSTM(32)),
29    layers.Dense(32, activation="relu"),
30    layers.Dense(len(label_names), activation="softmax"),
31])
32
33model.compile(
34    optimizer="adam",
35    loss="sparse_categorical_crossentropy",
36    metrics=["accuracy"],
37)
38
39model.fit(dataset, epochs=10, verbose=0)
40
41predictions = model.predict([
42    "minister discussed a new election strategy",
43    "the striker returned after a knee injury",
44], verbose=0)
45
46for row in predictions:
47    predicted_index = int(tf.argmax(row).numpy())
48    print(label_names[predicted_index], row.round(3))

This toy dataset is too small for a useful production model, but it demonstrates the full supervised workflow: labeled examples, vectorization, sequence model, training, and inference.

How the Pipeline Works

The TextVectorization layer builds a vocabulary and converts each document into integer token IDs. The embedding layer then learns a dense vector for each token. Similar words often end up with similar vectors after training because the network adjusts them to improve classification accuracy.

The bidirectional LSTM reads the sequence in both directions. That helps when important clues appear near either end of the sentence. For example, the meaning of a short support request can depend on the last few words just as much as the first few.

The final dense layer with softmax turns the learned document representation into class probabilities. In a binary problem you might use one output unit with sigmoid instead, but multi-class classification usually uses softmax.

Choosing a Model Family

For short texts and modest datasets, CNNs and LSTMs are still reasonable baselines. They train faster than large transformer models and are easier to deploy in constrained environments.

For higher accuracy on complex language tasks, pretrained transformers usually win. Fine-tuning a model such as BERT or DistilBERT can deliver strong results because the encoder already knows a lot about language before it sees your labels.

That said, starting with the biggest model is often a mistake. In many business applications, the hard part is not architecture selection but getting clean labels, balanced classes, and realistic evaluation data.

Evaluation Matters More Than Accuracy Alone

Accuracy can be misleading when classes are imbalanced. Suppose 90 percent of your documents belong to one category. A weak model can score 90 percent accuracy by always predicting that majority class.

For that reason, document classification is usually evaluated with:

  • Precision, to measure false positives
  • Recall, to measure false negatives
  • F1 score, to balance precision and recall
  • Confusion matrices, to inspect which classes are mixed up

You should also create a validation set that reflects production traffic. If training data contains long formal articles but production inputs are short, noisy user messages, the model may look good offline and fail in real use.

Common Pitfalls

The first pitfall is training on badly labeled data. Deep models can memorize label noise surprisingly well, which means they may look strong on the training set while learning the wrong patterns.

Another mistake is leaking information from validation into training. If you fit the vocabulary, normalize text, or duplicate near-identical documents across both splits, your evaluation becomes overly optimistic.

A third problem is ignoring class imbalance. If some categories are rare, use stratified splits, class weights, or targeted data collection instead of trusting raw accuracy.

Finally, do not overcomplicate the architecture before establishing a baseline. A small embedding model can tell you whether the labels and preprocessing are sensible. If that baseline fails, a larger network will usually fail more expensively.

Summary

  • Document classification is a supervised task that maps text to predefined labels.
  • Deep learning reduces manual feature engineering by learning representations from raw text.
  • A practical pipeline includes text vectorization, embeddings, a sequence or encoder model, and a classifier head.
  • Evaluation should include precision, recall, F1 score, and class-level error analysis.
  • Clean labels and realistic validation data usually matter more than choosing the fanciest model.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.