Toy Text Classifier

A small experimental project for fine-tuning and evaluating DistilBERT-based text classifiers, with experiment tracking via Weights & Biases.

This is a toy/learning project for text classification using Hugging Face transformers and PyTorch. It fine-tunes a distilbert-base-uncased-finetuned-sst-2-english model on the Rotten Tomatoes movie review dataset (binary sentiment classification), and logs training metrics to Weights & Biases (wandb). The project also includes a small pytest suite covering tokenization/preprocessing behavior.

• Editor: VS Code
• Python Version: 3.14 (see .python-version)
• Dependency Manager: uv (see pyproject.toml / uv.lock)

• General Purpose: accelerate
• Data Manipulation: datasets, numpy, pandas
• Machine Learning: torch, transformers
• Experiment Tracking: wandb
• Testing/Linting: pytest, ruff

Install dependencies with:

uv sync

• Rotten Tomatoes movie review dataset — the default dataset, loaded via the Hugging Face datasets library for training/evaluation.
• tests/data/ contains small CSV/JSONL fixture files (train.csv, test.csv, train.jsonl, test.jsonl) used both by the test suite and as example local --data-source inputs.

toy_text_classifier/dataset.py's load_data() accepts either a Hugging Face Hub dataset id (pulled at runtime via datasets.load_dataset(...)) or a local CSV/JSON(L) file/save_to_disk directory path, dispatching automatically based on whether the given source resolves to an existing local path.

Text is tokenized using the model's tokenizer (toy_text_classifier/features.py), with truncation and padding to a fixed max length, and dynamic padding at batch time via DataCollatorWithPadding (toy_text_classifier/modeling/train.py).

This project follows the cookiecutter-data-science layout, with functionality broken out into discrete, testable tools (dataset, features, modeling.train, modeling.predict) called by a CLI orchestrator.

├── data/{raw,interim,processed,external}/  # local data directories (gitignored contents)
├── models/                   # trained model checkpoints (gitignored contents)
├── notebooks/, references/, reports/figures/  # placeholders for future work
├── toy_text_classifier/
│   ├── config.py              # paths, default HF Hub ids, default hyperparameters
│   ├── dataset.py             # load_data(): HF Hub id or local path -> Dataset
│   ├── features.py            # tokenization/preprocessing
│   ├── modeling/
│   │   ├── common.py          # load_model_and_tokenizer(): HF Hub id or local path
│   │   ├── train.py           # train(): fine-tuning, optional wandb logging
│   │   └── predict.py         # predict(): inference on trained/base models
│   ├── plots.py                # placeholder (no visualization work done yet)
│   └── __main__.py             # CLI orchestrator (train/predict subcommands)
├── tests/
│   ├── conftest.py            # shared pytest fixtures (mocked tokenizer/model, CSV fixture paths)
│   ├── test_dataset.py, test_features.py, test_modeling_train.py,
│   │   test_modeling_predict.py, test_cli.py   # all run offline (external calls mocked)
│   └── data/                  # small CSV/JSONL fixtures used by tests
├── Makefile
├── pyproject.toml
└── uv.lock

Deliberately not included from the standard cookiecutter-data-science template: docs/ (no mkdocs skeleton — nothing to document in a toy project yet), requirements.txt and setup.cfg (this project fully commits to uv + pyproject.toml, and ruff reads its config from pyproject.toml, not a setup.cfg/flake8 file — adding either would just be an unmaintained duplicate).

make test        # or: uv run pytest

The suite runs fully offline — Hugging Face Hub downloads and wandb calls are mocked.

--data-source and --model-source (and the equivalents in config.py) each accept either a Hugging Face Hub id (e.g. cornell-movie-review-data/rotten_tomatoes, distilbert-base-uncased-finetuned-sst-2-english) or a local filesystem path (a CSV/JSON(L) file, a datasets.save_to_disk directory, or a locally saved model checkpoint directory).

make train
# or, with custom sources/hyperparameters:
uv run python -m toy_text_classifier train \
  --data-source cornell-movie-review-data/rotten_tomatoes \
  --model-source distilbert-base-uncased-finetuned-sst-2-english \
  --num-epochs 3 --no-wandb

Fine-tuning saves the resulting model + tokenizer to --output-dir (default models/trained/). There is no separate "download and cache the base model" step — from_pretrained handles Hub caching transparently.

make predict TEXT="this movie was great"
# or:
uv run python -m toy_text_classifier predict --model-path models/trained --text "great" --text "terrible"

Training runs are logged to Weights & Biases when --no-wandb is not passed (see the wandb/ run directories). No formal evaluation metrics or benchmark comparisons have been recorded yet — this is currently a work-in-progress experimentation project.

• Add held-out evaluation/metrics reporting (accuracy, F1) beyond training loss
• Implement plots.py (currently a placeholder; would need a matplotlib dependency)

• Base model: distilbert-base-uncased-finetuned-sst-2-english (Hugging Face)
• Dataset: Rotten Tomatoes movie review dataset (Cornell Movie Review Data)
• Experiment tracking: Weights & Biases

No license has been specified for this project yet.