Budget Tracker with Category Classifier

Budget Tracker with Category Classifier

Problem Statement

Manually categorizing bank transactions is one of the most tedious parts of personal finance management. Every month, you download a CSV export from your bank, and then spend 20+ minutes tagging each transaction as "Groceries," "Transport," "Utilities," etc. — just to understand where your money went. This manual process is error-prone and discouraging, which is why most people give up on budgeting entirely.

This script solves that problem by training a machine learning classifier on your historical categorized transactions, then automatically categorizing new transactions from any bank export. It persists the trained model to disk so you can reuse it month after month, and it gets smarter as you feed it more labeled data.

How It Works

The script operates in two modes: training and prediction. In training mode (--train), it reads a CSV containing both transaction descriptions and their known categories, builds a TF-IDF feature matrix from the text, and fits a LinearSVC classifier (fast, accurate, and handles high-dimensional sparse text well). The trained pipeline — vectorizer plus classifier — is serialized to disk with pickle.

In prediction mode, the script loads the persisted model, reads a new bank export, classifies every transaction description, and writes an enriched CSV with an added category column. Descriptions are normalized (lowercased, punctuation stripped, digits replaced with a token) before vectorization to reduce noise from transaction IDs and merchant codes. The script also reports a classification confidence proxy using the decision function margin, so you can flag low-confidence predictions for manual review. Design patterns used: a small Config dataclass for dependency injection, a Pipeline for reproducibility, and argparse subcommand-free flags for CLI simplicity.

Full Script

#!/usr/bin/env python3
"""
Budget Tracker with Category Classifier

Trains a text classifier on historical categorized transactions and uses it
to auto-categorize new bank CSV exports. Supports both training and prediction
modes from a single entry point.

Usage:
    python3 budget_track.py --train --input labeled.csv --model categories.pkl
    python3 budget_track.py --input bank_export.csv --model categories.pkl --output categorized.csv
"""

from __future__ import annotations

import argparse
import logging
import os
import pickle
import re
import sys
from dataclasses import dataclass
from pathlib import Path

import pandas as pd
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.model_selection import cross_val_score
from sklearn.pipeline import Pipeline
from sklearn.svm import LinearSVC

# ---------------------------------------------------------------------------
# Configuration (env vars with sensible defaults)
# ---------------------------------------------------------------------------
LOG_LEVEL = os.environ.get("BUDGET_LOG_LEVEL", "INFO").upper()
DEFAULT_TEXT_COLUMN = os.environ.get("BUDGET_TEXT_COLUMN", "description")
DEFAULT_LABEL_COLUMN = os.environ.get("BUDGET_LABEL_COLUMN", "category")
DEFAULT_AMOUNT_COLUMN = os.environ.get("BUDGET_AMOUNT_COLUMN", "amount")
DEFAULT_MODEL_PATH = os.environ.get("BUDGET_MODEL_PATH", "categories.pkl")
MIN_TRAINING_ROWS = int(os.environ.get("BUDGET_MIN_ROWS", "20"))
LOW_CONFIDENCE_THRESHOLD = float(os.environ.get("BUDGET_CONF_THRESHOLD", "0.35"))

logging.basicConfig(
    level=LOG_LEVEL,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
log = logging.getLogger("budget_track")


# ---------------------------------------------------------------------------
# Text normalization
# ---------------------------------------------------------------------------
_NON_ALNUM = re.compile(r"[^a-z0-9\s]")
_DIGITS = re.compile(r"\b\d+\b")
_WHITESPACE = re.compile(r"\s+")


def normalize_text(text: str) -> str:
    """Lowercase, strip punctuation, collapse digits into a single token."""
    if not isinstance(text, str):
        return ""
    text = text.lower()
    text = _NON_ALNUM.sub(" ", text)
    text = _DIGITS.sub(" num ", text)
    return _WHITESPACE.sub(" ", text).strip()


# ---------------------------------------------------------------------------
# Config dataclass
# ---------------------------------------------------------------------------
@dataclass
class Config:
    input_path: Path
    model_path: Path
    output_path: Path | None
    text_column: str
    label_column: str
    amount_column: str
    train: bool


# ---------------------------------------------------------------------------
# Model building
# ---------------------------------------------------------------------------
def build_pipeline() -> Pipeline:
    """Create a fresh TF-IDF + LinearSVC pipeline."""
    return Pipeline(
        [
            (
                "tfidf",
                TfidfVectorizer(
                    preprocessor=normalize_text,
                    ngram_range=(1, 2),
                    min_df=1,
                    sublinear_tf=True,
                ),
            ),
            ("clf", LinearSVC(C=1.0, max_iter=5000)),
        ]
    )


def load_dataframe(path: Path) -> pd.DataFrame:
    """Read a CSV with helpful error context."""
    if not path.exists():
        raise FileNotFoundError(f"Input file not found: {path}")
    try:
        df = pd.read_csv(path)
    except pd.errors.ParserError as exc:
        raise ValueError(f"Could not parse CSV at {path}: {exc}") from exc
    if df.empty:
        raise ValueError(f"Input CSV at {path} is empty")
    return df


# ---------------------------------------------------------------------------
# Training
# ---------------------------------------------------------------------------
def train_model(cfg: Config) -> None:
    df = load_dataframe(cfg.input_path)

    missing = {cfg.text_column, cfg.label_column} - set(df.columns)
    if missing:
        raise KeyError(
            f"Training CSV must contain columns {missing}. "
            f"Found: {list(df.columns)}"
        )

    df = df.dropna(subset=[cfg.text_column, cfg.label_column]).copy()
    df = df[df[cfg.text_column].astype(str).str.strip() != ""]

    if len(df) < MIN_TRAINING_ROWS:
        raise ValueError(
            f"Need at least {MIN_TRAINING_ROWS} labeled rows to train, "
            f"got {len(df)}."
        )

    X = df[cfg.text_column].astype(str).tolist()
    y = df[cfg.label_column].astype(str).tolist()

    log.info("Training on %d rows across %d categories", len(X), len(set(y)))

    pipeline = build_pipeline()
    # 3-fold CV for a quick sanity check (skip if too few samples per class)
    min_class_count = df[cfg.label_column].value_counts().min()
    if min_class_count >= 3:
        scores = cross_val_score(pipeline, X, y, cv=3, scoring="accuracy")
        log.info("Cross-val accuracy: %.3f (+/- %.3f)", scores.mean(), scores.std())
    else:
        log.warning("Some classes have <3 samples; skipping cross-validation")

    pipeline.fit(X, y)

    cfg.model_path.parent.mkdir(parents=True, exist_ok=True)
    with cfg.model_path.open("wb") as fh:
        pickle.dump(pipeline, fh)
    log.info("Model saved to %s", cfg.model_path)


# ---------------------------------------------------------------------------
# Prediction
# ---------------------------------------------------------------------------
def load_model(path: Path) -> Pipeline:
    if not path.exists():
        raise FileNotFoundError(
            f"Model file not found: {path}. Train one first with --train."
        )
    with path.open("rb") as fh:
        model = pickle.load(fh)
    if not isinstance(model, Pipeline):
        raise TypeError(f"Unexpected object in {path}: {type(model).__name__}")
    return model


def predict(cfg: Config) -> None:
    model = load_model(cfg.model_path)
    df = load_dataframe(cfg.input_path)

    if cfg.text_column not in df.columns:
        raise KeyError(
            f"Input CSV missing text column '{cfg.text_column}'. "
            f"Found: {list(df.columns)}"
        )

    texts = df[cfg.text_column].fillna("").astype(str)
    df["category"] = model.predict(texts)

    # Confidence proxy: distance from the decision boundary, normalized.
    # Larger |decision| = more confident. We convert to a 0-1-ish score.
    try:
        decisions = model.decision_function(texts)
        if decisions.ndim == 1:
            margins = decisions
        else:
            sorted_d = -1 * (decisions.max(axis=1) - decisions.min(axis=1))
            margins = sorted_d
        df["confidence"] = 1.0 / (1.0 + abs(margins))
    except Exception as exc:  # noqa: BLE001 - best-effort confidence
        log.debug("Could not compute confidence scores: %s", exc)
        df["confidence"] = 1.0

    low_conf = df[df["confidence"] < LOW_CONFIDENCE_THRESHOLD]
    if not low_conf.empty:
        log.warning(
            "%d transaction(s) below confidence threshold %.2f — review these.",
            len(low_conf),
            LOW_CONFIDENCE_THRESHOLD,
        )

    if cfg.amount_column in df.columns:
        summary = (
            df.groupby("category")[cfg.amount_column]
            .agg(["count", "sum"])
            .sort_values("sum")
        )
        log.info("Spend summary by category:\n%s", summary.to_string())

    out = cfg.output_path or Path("categorized.csv")
    df.to_csv(out, index=False)
    log.info("Wrote %d categorized rows to %s", len(df), out)


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------
def parse_args(argv: list[str] | None = None) -> Config:
    parser = argparse.ArgumentParser(
        description="Budget tracker with automatic transaction categorization.",
    )
    parser.add_argument("--input", required=True, type=Path, help="Input CSV path")
    parser.add_argument(
        "--model",
        type=Path,
        default=Path(DEFAULT_MODEL_PATH),
        help=f"Model file path (default: {DEFAULT_MODEL_PATH})",
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=None,
        help="Output CSV path (default: categorized.csv)",
    )
    parser.add_argument(
        "--train",
        action="store_true",
        help="Train mode: input CSV must contain category labels",
    )
    parser.add_argument(
        "--text-column",
        default=DEFAULT_TEXT_COLUMN,
        help=f"Column with transaction description (default: {DEFAULT_TEXT_COLUMN})",
    )
    parser.add_argument(
        "--label-column",
        default=DEFAULT_LABEL_COLUMN,
        help=f"Column with category labels, training only (default: {DEFAULT_LABEL_COLUMN})",
    )
    parser.add_argument(
        "--amount-column",
        default=DEFAULT_AMOUNT_COLUMN,
        help=f"Column with transaction amount (default: {DEFAULT_AMOUNT_COLUMN})",
    )
    args = parser.parse_args(argv)

    return Config(
        input_path=args.input,
        model_path=args.model,
        output_path=args.output,
        text_column=args.text_column,
        label_column=args.label_column,
        amount_column=args.amount_column,
        train=args.train,
    )


def main(argv: list[str] | None = None) -> int:
    try:
        cfg = parse_args(argv)
        if cfg.train:
            train_model(cfg)
        else:
            predict(cfg)
        return 0
    except (FileNotFoundError, ValueError, KeyError, TypeError) as exc:
        log.error("%s", exc)
        return 2
    except KeyboardInterrupt:
        log.warning("Interrupted by user")
        return 130
    except Exception as exc:  # noqa: BLE001 - top-level safety net
        log.exception("Unexpected error: %s", exc)
        return 1


if __name__ == "__main__":
    sys.exit(main())

Configuration

CLI flags (--input, --model, --output, --train, --text-column, --label-column, --amount-column) override the defaults for a single run.

Usage Examples

1. Train on a labeled history export:

$ python3 budget_track.py --train --input labeled_history.csv --model categories.pkl
2024-06-01 10:12:03 [INFO] Training on 312 rows across 9 categories
2024-06-01 10:12:03 [INFO] Cross-val accuracy: 0.914 (+/- 0.021)
2024-06-01 10:12:04 [INFO] Model saved to categories.pkl

2. Categorize a fresh bank export:

$ python3 budget_track.py --input bank_export.csv --model categories.pkl --output categorized.csv
2024-06-01 10:15:22 [INFO] Spend summary by category:
              count      sum
category
Groceries        14  -412.87
Transport         8  -128.40
Utilities         3  -210.00
...
2024-06-01 10:15:22 [INFO] Wrote 47 categorized rows to categorized.csv

3. Point at a non-standard schema via env vars:

$ BUDGET_TEXT_COLUMN=memo BUDGET_AMOUNT_COLUMN=value \
    python3 budget_track.py --input chase.csv --model chase.pkl --output out.csv

Error Handling

• Missing input file → FileNotFoundError raised with the resolved path; exit code 2.
• Malformed CSV → pandas.errors.ParserError wrapped in a ValueError with the source path.
• Missing required columns → KeyError listing both the missing names and the columns actually present.
• Too few training rows → ValueError before any model fitting, preventing garbage models.
• Missing model at prediction time → FileNotFoundError with a hint to run --train first.
• Corrupt or wrong-typed pickle → TypeError if the unpickled object is not a Pipeline.
• Interrupt (Ctrl-C) → clean exit with code 130.
• Anything else → full traceback logged, exit code 1, so cron jobs and CI can distinguish failures.

Extending Ideas

• Active learning loop: After prediction, read a reviewed.csv where a human corrected low-confidence rows, and append it to the training set on the next --train run.
• Swap the classifier: Replace LinearSVC with LogisticRegression to get calibrated predict_proba probabilities instead of the margin-based confidence proxy.
• Rule-based overrides: Add a small rules.yaml (e.g., "AMAZON" -> Shopping) applied before the ML model for known merchants — ML is best used as a fallback.
• Persist per-user models: Namespace categories.pkl by account ID (categories_{account}.pkl) so different bank accounts with different spending patterns get tailored classifiers.
• Streaming ingestion: Wrap the prediction path in a function that accepts a DataFrame so it can be called from an Airflow task or a webhook handler instead of only from the CLI.