Two-Tower Movie Retrieval — MovieLens 100K

Candidate-retrieval system that takes a user_id and returns Top-K movie candidates, built end-to-end with production-style engineering discipline: temporal evaluation, leakage audits, a popularity baseline, bias/coverage guardrails, and an exportable, version-stamped serving index with a cold-start fallback.

Python 3.11–3.13 · TensorFlow 2.20 / Keras 3 · no TFRS · 101 tests · CI on 3.11 + 3.12 · Colab-verified · MIT

Portfolio project — demonstrates the retrieval stage of a modern recommender (retrieval → ranking → re-ranking), the stage that narrows a large catalogue to a few hundred candidates under tight latency budgets.

Results at a glance

Temporal test set · 943 users · one held-out future interaction each · full-catalogue scoring

metric @10	popularity (R0)	two-tower (R1)
Recall@10	0.066	0.058
NDCG@10	0.033	0.027
Recall@50	0.183	0.231
Catalogue coverage@10	5.4%	89.3%
Top-10%-popular share of slots	100%	6.3%
Gini exposure	0.99	0.48

Honest headline: the popularity baseline edges out the ID-only two-tower model at K=10 — a well-known result on small temporal datasets — while the two-tower model wins at K=50 (the regime that matters for a retrieval stage feeding a ranker) and delivers ~16× higher catalogue coverage with drastically lower popularity concentration. Slice analysis shows the model is strongest exactly where popularity fails: long-tail test items (2.6× head recall) and low-history users. Conclusions are stable under the positive-interaction-rule sensitivity check (rating ≥ 4 vs all ratings). Serving latency: p95 ≈ 1 ms per single-user query on a laptop CPU (the exact figure moves with machine load — artifacts/metrics.json carries the value from the run reported here).

Left: share of the catalogue that ever appears in a Top-10 list. Right: how concentrated exposure is — the popularity curve stays flat along the axis and then shoots up at the very end (a handful of movies take essentially every slot, Gini 0.99), while the two-tower curve sits much closer to the equality diagonal (Gini 0.48). This is the trade that the K=10 number alone hides, and the reason both are tracked as first-class metrics.

Everything above is reproduced end-to-end by the notebook — re-run on a clean Google Colab runtime (Python 3.12 / TF 2.20, 21/21 cells, ~5 min), which returns the same metrics to every digit reported here.

Problem

Platforms with large catalogues cannot score every item with an expensive ranking model on every request. A retrieval layer must return relevant candidates fast, and its offline evaluation must mirror serving reality: predicting future interactions from past ones. A random train/test split leaks the future and inflates every metric — so everything here is evaluated with a temporal leave-last-1-out split per user, enforced by an automated leakage audit that fails the pipeline on violation.

Left: five users' timelines — validation and test markers never sit to the left of train. Right: the same check across all 943 users, and the honest detail it exposes — for 422 users the validation interaction shares the same second as their last training interaction (415 for test vs validation), because MovieLens timestamps have one-second resolution and people rate several movies in one sitting. So the invariant the audit enforces is max(train) ≤ val ≤ test with a deterministic movie_id tie-break, not a strict inequality. Zero users violate it; writing "val/test always come after train" would have been an overclaim the data does not support.

Approach

ratings ──► data contract ──► temporal split (audited) ──► popularity baseline (R0)
   (SHA-256 verified download)        │
                                      └──► two-tower model (R1/R2) ──► full-catalogue eval
                                                    │                   Recall/NDCG/coverage/bias/slices
                                                    └──► brute-force SavedModel index
                                                         + seen-filter + popularity fallback

Model: two-tower user_id/movie_id embeddings (dim 32), dot-product affinity, in-batch sampled softmax with accidental-hit masking — implemented in ~60 lines of pure TensorFlow/Keras 3 (TFRS is in maintenance mode and Keras-3-incompatible; owning the loss keeps the dependency surface small and the math auditable).
Selection discipline: R1 (dim 32) vs R2 (dim 64) compared on validation Recall@10 only; the test set is scored exactly once, after all decisions.
Serving: pure-TF SavedModel index (no Keras dependency), reload-consistency checked at export time, input validation, seen-item filtering, and model_version/index_version stamped into every response.

Where this sits in a production recommender

flowchart LR
    ES[Event stream<br/>clicks / watches / ratings] --> FS[(Feature &<br/>history store)]
    FS --> UT[User tower<br/>→ query vector]
    FS -. nightly retrain .-> TT[Two-tower training]
    TT --> IT[Item tower] --> IDX[(ANN candidate index<br/>MVP: BruteForce SavedModel)]
    UT --> IDX
    IDX --> F[Seen-item &<br/>business/safety filters]
    F --> R[Ranking stage<br/>out of MVP scope]
    R --> UI[UI exposure]
    UI --> FB[Feedback / labels] --> ES

    classDef mvp fill:#dbeafe,stroke:#2563eb;
    classDef out fill:#f3f4f6,stroke:#9ca3af,stroke-dasharray:4 3;
    class UT,IT,IDX,F,TT mvp
    class R,UI out

Blue = built here · dashed grey = designed but outside the 10-hour scope. The arrow from exposure back to the event stream is why coverage and popularity bias are tracked as first-class metrics: today's recommendations become tomorrow's training labels.

Quickstart

python3.11 -m venv .venv && source .venv/bin/activate
pip install -e ".[dev]"                # package + CLI + test/notebook tooling

movie-retrieval all --sensitivity      # download → validate → split → train → evaluate → export
movie-retrieval recommend --user-id 42 --k 10
pytest                                 # 101 tests incl. leakage & zip-slip guards
pytest --cov --cov-fail-under=65       # the coverage gate CI enforces on every push

Three dependency files, three different jobs — complements, not duplicates:

file	pins	use it when
`pyproject.toml`	version ranges + the `movie-retrieval` console script	installing the project itself (recommended)
`requirements.txt`	the same runtime ranges plus notebook deps, no package install	you only need the libraries — a bare Colab/CI runtime
`requirements-lock.txt`	exact resolved versions of the run reported above	reproducing the published metrics bit-for-bit

The dataset (~5 MB) is downloaded at runtime from GroupLens and verified against a pinned SHA-256. It is never committed (MovieLens research-use terms prohibit redistribution).

Try the retrieval index yourself

pip install -e ".[demo]"               # adds Gradio; the library and CLI don't need it
python app.py                          # http://127.0.0.1:7860

A small UI over the exported index: type a user id, get the Top-K candidates a ranking stage would receive. It calls the same RetrievalService as the CLI on the same artifacts, so the input validation, seen-item filtering and cold-start fallback are the real ones — enter an id outside the training vocabulary (brand-new-user) and the response switches to the popularity fallback in front of you. Requires movie-retrieval all to have produced the artifacts first; nothing is uploaded and no dataset is bundled.

Repository layout

├── src/movie_retrieval/      # config, data, splits, baseline, model, evaluate, index, pipeline, cli, demo
├── app.py                    # entry point for the Gradio demo over the exported index
├── tests/                    # unit + integration tests (leakage, metrics, index reload, security)
├── notebooks/movielens_two_tower_retrieval.ipynb   # 19-section narrative notebook (executed)
│   └── ..._Colab_Ran.ipynb   # same notebook, executed top-to-bottom on a clean Colab runtime
├── docs/img/                 # plots exported from the notebook for this README
├── .github/workflows/ci.yml  # ruff + pytest on Python 3.11 / 3.12
├── artifacts/                # generated: model, index, vocab, metrics (gitignored)
├── requirements.txt          # runtime + notebook dependency ranges
└── requirements-lock.txt     # frozen dependency versions (exact reproduction)

Key engineering decisions

Temporal leave-last-k split with hard audit — LeakageError stops the run on any future leakage, missing history, or train/test overlap; split rule is versioned in artifacts/split_config.json.
Train-only statistics — vocabularies and popularity counts never see val/test; OOV test items are reported (0.2%), not silently dropped.
SUM loss reduction (TFRS semantics) — with MEAN reduction, per-parameter gradients are ~batch-size smaller and Adagrad stalls; a subtle bug that presents as "flat loss".
Full-catalogue evaluation — 1,682 items makes exact scoring cheap; no sampled-negative bias in reported metrics.
Fail-hard artifact export — export aborts unless the reloaded index reproduces the in-memory Top-K exactly.
Security hygiene — HTTPS + pinned SHA-256 download, zip-slip-safe extraction, input validation at the serving API, no secrets/data in the repo.

Model card


Model	`ml100k-retrieval-v1` — two-tower, dim-32 ID embeddings, dot-product affinity, in-batch sampled softmax (SUM reduction) with accidental-hit masking
Training	Adagrad lr 0.1, 15 epochs, batch 256, seed 42 · TensorFlow 2.20 / Keras 3, no TFRS
Selection	R1 (dim 32) over R2 (dim 64 + L2) on validation Recall@10; the test set is scored once, after every decision
Data	MovieLens 100K — 100,000 ratings, 943 users, 1,682 movies, collected 1997–1998; every rating counts as an interaction (`rating ≥ 4` reported as a sensitivity check)
Intended use	Portfolio and education — demonstrating retrieval-stage methodology
Out of scope	Commercial or production use (prohibited by the MovieLens terms), ranking, and any setting where recommendations could cause harm without human curation
Cold start	Unknown user → popularity Top-K with `fallback_used: true` in the response; unknown movie → not retrievable at all, which is the limitation a content tower would fix
Known risks	Retrieval trained on logged interactions amplifies exposure bias — the reason coverage and Gini are tracked as first-class metrics; there is no content-safety layer; behavioural data is personal data even when de-identified

Limitations

No impression data: unobserved ≠ disliked; offline Recall is a proxy — only an online A/B test measures real lift.
1998-era, 100K-interaction dataset: methodology transfers, taste conclusions don't.
One test interaction per user → K=10 differences of ±0.01 are noise-level.
ID-only towers cannot embed new movies (item cold-start) — needs content features (title/genre tower), listed as the first stretch goal.

License & data terms

Code: MIT (LICENSE) — covers this repository's source only.

Dataset: MovieLens 100K by GroupLens, and the MIT license does not extend to it: research use only, no redistribution, no commercial use. It is downloaded at runtime against a pinned SHA-256 and excluded from version control, so nothing in this repository redistributes it.

FilesExpand file tree

README.md

Latest commit

History