Machine Learning for Document Classification in Python

Your shared drive has thousands of files, and someone has to decide which ones are invoices, contracts or support tickets.

Quick answer: machine learning for document classification turns each document into numbers (TF-IDF word weights) and trains a classifier on labeled examples, so it can sort new files on its own. In scikit-learn, that’s a two-step pipeline:

model = make_pipeline(TfidfVectorizer(), ComplementNB())
model.fit(train_texts, train_labels)
model.predict(["My 2019 Honda Accord stalls at red lights"])

In this tutorial, we’ll build machine learning for document classification step by step: load labeled documents, extract features, train and compare five models, read the scores, then classify and save.

What You Need Before You Start Classifying Documents

We ran every example on Python 3.12.5 with scikit-learn 1.9.1. Here’s the short list:

• Python 3.10 or later and pip install scikit-learn joblib.
• About 15 MB of disk space; scikit-learn downloads the 20 Newsgroups dataset on the first run and caches it.
• A terminal. Our screenshots come from Command Prompt on Windows.

The running sample is 20 Newsgroups, a classic labeled set of real forum posts. We keep four topics, graphics, autos, medicine and space, so each one behaves like a document type in a company archive.

What is document classification in machine learning?

Document classification means assigning a document to one or more predefined categories. A model learns from documents someone has already labeled, which makes it supervised learning.

There are three flavors: binary sorts into two groups, like spam or not spam. Multi-class picks exactly one of several labels. Multi-label lets one document carry several tags, like an invoice that’s also a contract.

Scanned PDFs and images need one extra step first: OCR (optical character recognition) to pull out the text. After that, the pipeline is the same. For the wider NLP picture, see our guide to machine learning techniques for text.

Step 1: Load the Labeled Documents

fetch_20newsgroups() returns the texts, the numeric labels and the label names. We strip headers, footers and quoted replies, because they leak the answer (more on that in the mistakes section):

from sklearn.datasets import fetch_20newsgroups

CATS = ["comp.graphics", "rec.autos", "sci.med", "sci.space"]
KW = dict(categories=CATS, remove=("headers", "footers", "quotes"),
          shuffle=True, random_state=42)
train = fetch_20newsgroups(subset="train", **KW)
test = fetch_20newsgroups(subset="test", **KW)

from collections import Counter

print(f"Training documents: {len(train.data):,}")
print(f"Test documents:     {len(test.data):,}")
counts = Counter(train.target)
for i, name in enumerate(train.target_names):
    print(f"  {name:<14} {counts[i]:>4} training docs")

empty = sum(1 for d in train.data if not d.strip())
print(f"Empty after cleanup: {empty}")

doc = " ".join(train.data[1].split())[:110]
print(f"\nDocument 1 ({train.target_names[train.target[1]]}):")
print(doc + "...")

Output:

Training documents: 2,365
Test documents:     1,575
  comp.graphics   584 training docs
  rec.autos       594 training docs
  sci.med         594 training docs
  sci.space       593 training docs
Empty after cleanup: 77

Document 1 (rec.autos):
Could be due to the rear-end ratio also. Usually automatics have different rear-ends than manuals, from my lim...

So we have 2,365 training and 1,575 test documents. The classes are balanced, and 77 documents ended up empty after the cleanup, which is normal for forum data. Cleaning text like this is part of data preprocessing in machine learning.

Note: for your own files, load a CSV with text and label columns. Our guide to reading a CSV file with pandas covers that part.

Step 2: Turn Documents into TF-IDF Features

Models need numbers, so TF-IDF (term frequency times inverse document frequency) gives a word a high weight when it’s frequent in one document but rare across the collection:

from sklearn.datasets import fetch_20newsgroups

CATS = ["comp.graphics", "rec.autos", "sci.med", "sci.space"]
KW = dict(categories=CATS, remove=("headers", "footers", "quotes"),
          shuffle=True, random_state=42)
train = fetch_20newsgroups(subset="train", **KW)
test = fetch_20newsgroups(subset="test", **KW)

from sklearn.feature_extraction.text import TfidfVectorizer

vec = TfidfVectorizer(stop_words="english", sublinear_tf=True,
                      min_df=2, max_df=0.5)
X_train = vec.fit_transform(train.data)   # learn vocabulary on train
X_test = vec.transform(test.data)         # reuse it on test

print("Train matrix:", X_train.shape)
print("Test matrix: ", X_test.shape)
print(f"Non-zero cells: {X_train.nnz:,}")

terms = vec.get_feature_names_out()
row = X_train[1].toarray().ravel()
print("Top terms in document 1:")
for i in row.argsort()[::-1][:6]:
    print(f"  {terms[i]:<12} {row[i]:.3f}")

Output:

Train matrix: (2365, 13250)
Test matrix:  (1575, 13250)
Non-zero cells: 131,536
Top terms in document 1:
  rear         0.445
  anyways      0.351
  automatics   0.334
  manuals      0.322
  ends         0.317
  ratio        0.288

Each document is now a row with 13,250 columns, one per word. The top terms for the car post, like rear, automatics and manuals, show why this works.

We fit the vectorizer on the training set only and reuse it on the test set. stop_words="english" drops filler words, min_df=2 ignores one-off typos and max_df=0.5 drops words found in over half the documents. More options live in our feature extraction in machine learning guide.

Step 3: Train a Naive Bayes Document Classifier

A pipeline glues the vectorizer and the model together, so the same steps run at training and prediction time. Multinomial Naive Bayes is a fast baseline that works on word weights:

from sklearn.datasets import fetch_20newsgroups

CATS = ["comp.graphics", "rec.autos", "sci.med", "sci.space"]
KW = dict(categories=CATS, remove=("headers", "footers", "quotes"),
          shuffle=True, random_state=42)
train = fetch_20newsgroups(subset="train", **KW)
test = fetch_20newsgroups(subset="test", **KW)

from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.naive_bayes import MultinomialNB
from sklearn.pipeline import make_pipeline
from sklearn.metrics import accuracy_score

model = make_pipeline(
    TfidfVectorizer(stop_words="english", sublinear_tf=True,
                    min_df=2, max_df=0.5),
    MultinomialNB(alpha=0.1),
)
model.fit(train.data, train.target)
pred = model.predict(test.data)
print(f"Naive Bayes accuracy: {accuracy_score(test.target, pred):.3f}")

Output:

Naive Bayes accuracy: 0.890

That’s 89.0% of 1,575 unseen documents labeled correctly. Accuracy alone can hide weak classes, though. Our scikit-learn accuracy_score guide explains when it misleads.

Step 4: Compare Five Document Classification Models

Next we swap the classifier and keep the same features. We tested both Naive Bayes variants, logistic regression in scikit-learn, a linear support vector machine and a random forest:

from sklearn.datasets import fetch_20newsgroups

CATS = ["comp.graphics", "rec.autos", "sci.med", "sci.space"]
KW = dict(categories=CATS, remove=("headers", "footers", "quotes"),
          shuffle=True, random_state=42)
train = fetch_20newsgroups(subset="train", **KW)
test = fetch_20newsgroups(subset="test", **KW)

from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.naive_bayes import MultinomialNB, ComplementNB
from sklearn.linear_model import LogisticRegression
from sklearn.svm import LinearSVC
from sklearn.ensemble import RandomForestClassifier
from sklearn.pipeline import make_pipeline
from sklearn.metrics import accuracy_score, f1_score

models = {
    "MultinomialNB": MultinomialNB(alpha=0.1),
    "ComplementNB": ComplementNB(alpha=0.3),
    "LogisticRegression": LogisticRegression(C=10, max_iter=1000),
    "LinearSVC": LinearSVC(C=0.5),
    "RandomForest": RandomForestClassifier(n_estimators=300,
                                           random_state=42),
}
print(f"{'Model':<20}{'Accuracy':>9}{'Macro F1':>10}")
for name, clf in models.items():
    pipe = make_pipeline(
        TfidfVectorizer(stop_words="english", sublinear_tf=True,
                        min_df=2, max_df=0.5), clf)
    pipe.fit(train.data, train.target)
    pred = pipe.predict(test.data)
    acc = accuracy_score(test.target, pred)
    f1 = f1_score(test.target, pred, average="macro")
    print(f"{name:<20}{acc:>9.3f}{f1:>10.3f}")

Output:

Model                Accuracy  Macro F1
MultinomialNB           0.890     0.890
ComplementNB            0.890     0.891
LogisticRegression      0.863     0.864
LinearSVC               0.864     0.865
RandomForest            0.803     0.804

This surprised us a bit. Most guides crown the linear SVM, but on 2,365 short documents ComplementNB had the best macro F1 at 0.891. LinearSVC scored 0.864 and the random forest 0.803.

Pro Tip: rerun this comparison on your own data. With tens of thousands of documents, linear models usually catch up.

Step 5: Evaluate the Classifier with a Report and Confusion Matrix

Now we take the winner, ComplementNB, and look at each class. Precision answers “when it says sci.med, how often is it right?” Recall answers “how many sci.med posts did it find?”:

from sklearn.datasets import fetch_20newsgroups

CATS = ["comp.graphics", "rec.autos", "sci.med", "sci.space"]
KW = dict(categories=CATS, remove=("headers", "footers", "quotes"),
          shuffle=True, random_state=42)
train = fetch_20newsgroups(subset="train", **KW)
test = fetch_20newsgroups(subset="test", **KW)

from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.naive_bayes import ComplementNB
from sklearn.pipeline import make_pipeline
from sklearn.metrics import classification_report, confusion_matrix

model = make_pipeline(
    TfidfVectorizer(stop_words="english", sublinear_tf=True,
                    min_df=2, max_df=0.5),
    ComplementNB(alpha=0.3),
)
model.fit(train.data, train.target)
pred = model.predict(test.data)

print(classification_report(test.target, pred, digits=3,
                            target_names=train.target_names))
print("Confusion matrix (rows = true, columns = predicted):")
print(confusion_matrix(test.target, pred))

Output:

               precision    recall  f1-score   support

comp.graphics      0.796     0.951     0.867       389
    rec.autos      0.902     0.884     0.893       396
      sci.med      0.953     0.876     0.913       396
    sci.space      0.936     0.850     0.891       394

     accuracy                          0.890      1575
    macro avg      0.897     0.890     0.891      1575
 weighted avg      0.897     0.890     0.891      1575

Confusion matrix (rows = true, columns = predicted):
[[370  10   2   7]
 [ 32 350   5   9]
 [ 29  13 347   7]
 [ 34  15  10 335]]

Most errors land in the comp.graphics column: 32 auto posts, 29 medical posts and 34 space posts were filed there. That’s why comp.graphics has the highest recall (0.951) but the lowest precision (0.796). Our scikit-learn confusion matrix tutorial shows how to plot it.

Step 6: Classify New Documents and Save the Model

We save the pipeline with joblib, load it back like a separate app would, and score four new 2026 support questions:

from sklearn.datasets import fetch_20newsgroups

CATS = ["comp.graphics", "rec.autos", "sci.med", "sci.space"]
KW = dict(categories=CATS, remove=("headers", "footers", "quotes"),
          shuffle=True, random_state=42)
train = fetch_20newsgroups(subset="train", **KW)
test = fetch_20newsgroups(subset="test", **KW)

import joblib
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.naive_bayes import ComplementNB
from sklearn.pipeline import make_pipeline

model = make_pipeline(
    TfidfVectorizer(stop_words="english", sublinear_tf=True,
                    min_df=2, max_df=0.5),
    ComplementNB(alpha=0.3),
)
model.fit(train.data, train.target)
joblib.dump(model, "doc_classifier.joblib")

loaded = joblib.load("doc_classifier.joblib")
new_docs = [
    "My 2019 Honda Accord stalls at red lights in Phoenix. Is it "
    "the idle air valve or the transmission?",
    "NASA engineers are testing the Orion heat shield and the "
    "orbit plan before the next Artemis launch.",
    "Which GPU renders 3D models faster in Blender, and does "
    "OpenGL matter for ray tracing?",
    "My doctor in Boston says the new flu vaccine is safe with "
    "my blood pressure medication.",
]
probs = loaded.predict_proba(new_docs)
for doc, p in zip(new_docs, probs):
    best = p.argmax()
    label = train.target_names[best]
    print(f"{label:<14} {p[best]:6.1%}  {doc[:45]}...")

Output:

rec.autos       91.5%  My 2019 Honda Accord stalls at red lights in ...
sci.space       85.3%  NASA engineers are testing the Orion heat shi...
comp.graphics   80.1%  Which GPU renders 3D models faster in Blender...
sci.med         90.0%  My doctor in Boston says the new flu vaccine ...

All four landed in the right category, with confidence between 80.1% and 91.5%. We’d route low-confidence calls to a person. To serve the model over HTTP, see how to deploy a machine learning model as a REST API.

Document classification algorithms in scikit-learn: cheat sheet

Here’s what we measured on the same test set:

Model	Accuracy	Macro F1	Best for
ComplementNB	0.890	0.891	Small or imbalanced text sets
MultinomialNB	0.890	0.890	Fast baseline, word counts
LinearSVC	0.864	0.865	Large sparse data, no probabilities
LogisticRegression	0.863	0.864	Calibrated probabilities
RandomForest	0.803	0.804	Mixed features, not raw text

How do you classify documents with more than one label?

A late-fee notice can be an invoice and a contract matter at once. For that, we turn the tags into a 0/1 matrix with MultiLabelBinarizer and train one classifier per tag:

from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.multiclass import OneVsRestClassifier
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import MultiLabelBinarizer

docs = [
    "Invoice 2026-114 for $4,250.00 is due March 15, 2026.",
    "Please pay the attached invoice by Friday.",
    "This agreement renews on July 1, 2026 unless canceled.",
    "Contract amendment: payment terms change to net 45.",
    "Invoice past due: late fee of $75.00 applies per contract.",
    "Signed lease agreement for the Denver office.",
    "Reminder: invoice 2026-201 balance is $980.50.",
    "The vendor contract includes an invoice schedule.",
]
tags = [["invoice"], ["invoice"], ["contract"], ["contract"],
        ["invoice", "contract"], ["contract"], ["invoice"],
        ["invoice", "contract"]]

mlb = MultiLabelBinarizer()
Y = mlb.fit_transform(tags)
model = make_pipeline(
    TfidfVectorizer(),
    OneVsRestClassifier(LogisticRegression(C=10)),
)
model.fit(docs, Y)

new = ["Contract requires each invoice within 30 days.",
       "Invoice 2026-310 total: $1,120.00."]
for doc, row in zip(new, model.predict(new)):
    print(mlb.classes_[row == 1].tolist(), "<-", doc)

Output:

['contract', 'invoice'] <- Contract requires each invoice within 30 days.
['invoice'] <- Invoice 2026-310 total: $1,120.00.

The first new document got both tags and the second only invoice. For a deep-learning version, read about large-scale multi-label text classification with Keras.

Common Mistakes to Avoid in Document Classification

The biggest one we reproduced is leakage: keeping text that gives the answer away. Newsgroup headers include the group name and the authors’ email signatures, so the model learns those instead of the topic:

from sklearn.datasets import fetch_20newsgroups
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.svm import LinearSVC
from sklearn.pipeline import make_pipeline

CATS = ["comp.graphics", "rec.autos", "sci.med", "sci.space"]
for remove in [(), ("headers", "footers", "quotes")]:
    kw = dict(categories=CATS, remove=remove)
    train = fetch_20newsgroups(subset="train", **kw)
    test = fetch_20newsgroups(subset="test", **kw)
    model = make_pipeline(TfidfVectorizer(sublinear_tf=True),
                          LinearSVC(C=0.5))
    model.fit(train.data, train.target)
    acc = model.score(test.data, test.target)
    print(f"remove={remove or 'nothing'}: accuracy {acc:.3f}")

Output:

remove=nothing: accuracy 0.947
remove=('headers', 'footers', 'quotes'): accuracy 0.863

With headers left in, LinearSVC looked like a 94.7% model. On clean text, it’s really 86.3%. Your own files leak the same way through file names, folder paths, email subjects or invoice numbers that contain the label.

Also fit the vectorizer on the training split only, and never score a model on documents it trained on.

Should you use deep learning for document classification?

For a few thousand documents and clear categories, the TF-IDF pipeline above is hard to beat on speed and cost. Transformer models pay off with subtle categories, longer documents or many languages. Our comparison of machine learning vs neural networks covers that trade-off.

If you want to try that route, start with our text classification with a Keras transformer build and compare it against these numbers.

Next Steps after building a document classifier

Swap in your own labeled documents and rerun Steps 4 and 5. To explain predictions to stakeholders, our guide to interpretable machine learning with Python is the natural next read.

We’d keep these machine learning guides open while you build:

• Python libraries for machine learning
• The machine learning life cycle
• Text classification using FNet in Keras
• Machine learning image recognition

Frequently asked questions