Spaces:
Sleeping
Sleeping
Commit ·
d1ecac0
0
Parent(s):
deploy: 17ac72a6e2a5a8c6e90ff79748114d9b6dd6aa91 with Space config header
Browse filesA single-commit snapshot of the Space build context (88 files) of
https://github.com/apoorvrajdev/image-captioning-system/commit/17ac72a6e2a5a8c6e90ff79748114d9b6dd6aa91, with no GitHub history (ADR-030).
Source-Commit: 17ac72a6e2a5a8c6e90ff79748114d9b6dd6aa91
Deploy-Run: https://github.com/apoorvrajdev/image-captioning-system/actions/runs/39059105674
This view is limited to 50 files because it contains too many changes. See raw diff
- .dockerignore +71 -0
- .gitattributes +39 -0
- Dockerfile +58 -0
- README.md +391 -0
- backend/app/__init__.py +1 -0
- backend/app/api/__init__.py +5 -0
- backend/app/api/routes.py +170 -0
- backend/app/core/__init__.py +18 -0
- backend/app/core/body_limit.py +71 -0
- backend/app/core/config.py +120 -0
- backend/app/core/logging.py +87 -0
- backend/app/main.py +110 -0
- backend/app/schemas/__init__.py +5 -0
- backend/app/schemas/caption.py +65 -0
- backend/app/services/__init__.py +5 -0
- backend/app/services/predictor_service.py +145 -0
- backend/app/services/weights_loader.py +257 -0
- backend/app/tests/__init__.py +0 -0
- backend/app/tests/conftest.py +101 -0
- backend/app/tests/image_bytes.py +47 -0
- backend/app/tests/test_backend_settings.py +34 -0
- backend/app/tests/test_body_size_limit.py +207 -0
- backend/app/tests/test_captions.py +231 -0
- backend/app/tests/test_health.py +56 -0
- backend/app/tests/test_image_guard.py +158 -0
- backend/app/tests/test_inference_admission.py +325 -0
- backend/app/tests/test_startup_integrity.py +154 -0
- backend/app/tests/test_weights_loader.py +400 -0
- backend/app/utils/__init__.py +5 -0
- backend/app/utils/image.py +147 -0
- configs/base.yaml +78 -0
- configs/train/debug.yaml +18 -0
- configs/train/stabilized.yaml +84 -0
- models/v1.0.0/vocab.json +54 -0
- pyproject.toml +248 -0
- requirements.txt +51 -0
- src/captioning/__init__.py +23 -0
- src/captioning/baselines/__init__.py +24 -0
- src/captioning/baselines/base.py +64 -0
- src/captioning/baselines/cnn.py +101 -0
- src/captioning/baselines/hf.py +126 -0
- src/captioning/config/__init__.py +30 -0
- src/captioning/config/loader.py +73 -0
- src/captioning/config/schema.py +289 -0
- src/captioning/data/__init__.py +21 -0
- src/captioning/data/coco.py +92 -0
- src/captioning/data/pipeline.py +92 -0
- src/captioning/data/splits.py +86 -0
- src/captioning/evaluation/__init__.py +90 -0
- src/captioning/evaluation/benchmark.py +128 -0
.dockerignore
ADDED
|
@@ -0,0 +1,71 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# =============================================================================
|
| 2 |
+
# .dockerignore — keep the build context tiny.
|
| 3 |
+
# Everything not needed by the backend runtime is excluded.
|
| 4 |
+
# =============================================================================
|
| 5 |
+
|
| 6 |
+
# --- VCS / IDE ---------------------------------------------------------------
|
| 7 |
+
.git
|
| 8 |
+
.github
|
| 9 |
+
.gitignore
|
| 10 |
+
.gitattributes
|
| 11 |
+
.vscode
|
| 12 |
+
.idea
|
| 13 |
+
*.swp
|
| 14 |
+
|
| 15 |
+
# --- Python caches & build artefacts -----------------------------------------
|
| 16 |
+
__pycache__
|
| 17 |
+
*.py[cod]
|
| 18 |
+
*$py.class
|
| 19 |
+
*.egg-info
|
| 20 |
+
.eggs
|
| 21 |
+
build
|
| 22 |
+
dist
|
| 23 |
+
.pytest_cache
|
| 24 |
+
.mypy_cache
|
| 25 |
+
.ruff_cache
|
| 26 |
+
.coverage
|
| 27 |
+
.coverage.*
|
| 28 |
+
htmlcov
|
| 29 |
+
.tox
|
| 30 |
+
|
| 31 |
+
# --- Virtual envs ------------------------------------------------------------
|
| 32 |
+
.venv
|
| 33 |
+
venv
|
| 34 |
+
env
|
| 35 |
+
|
| 36 |
+
# --- Frontend (deployed separately to Vercel) --------------------------------
|
| 37 |
+
frontend
|
| 38 |
+
node_modules
|
| 39 |
+
|
| 40 |
+
# --- Tests, notebooks, docs --------------------------------------------------
|
| 41 |
+
tests
|
| 42 |
+
notebooks
|
| 43 |
+
docs
|
| 44 |
+
assets
|
| 45 |
+
scripts
|
| 46 |
+
|
| 47 |
+
# --- Local dev / secrets -----------------------------------------------------
|
| 48 |
+
.env
|
| 49 |
+
.env.*
|
| 50 |
+
!.env.example
|
| 51 |
+
*.local
|
| 52 |
+
|
| 53 |
+
# --- OS junk -----------------------------------------------------------------
|
| 54 |
+
.DS_Store
|
| 55 |
+
Thumbs.db
|
| 56 |
+
|
| 57 |
+
# --- Docker / CI configs (not needed inside the image) -----------------------
|
| 58 |
+
Dockerfile
|
| 59 |
+
.dockerignore
|
| 60 |
+
.pre-commit-config.yaml
|
| 61 |
+
Makefile
|
| 62 |
+
CLAUDE.md
|
| 63 |
+
|
| 64 |
+
# --- Lockfiles / dev-only deps -----------------------------------------------
|
| 65 |
+
requirements-dev.txt
|
| 66 |
+
requirements-eval.txt
|
| 67 |
+
|
| 68 |
+
# --- Model artefacts: keep ONLY the production version -----------------------
|
| 69 |
+
# (Add explicit exceptions here if you publish multiple versions.)
|
| 70 |
+
models/**/*.tar
|
| 71 |
+
models/**/*.zip
|
.gitattributes
ADDED
|
@@ -0,0 +1,39 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
*.7z filter=lfs diff=lfs merge=lfs -text
|
| 2 |
+
*.arrow filter=lfs diff=lfs merge=lfs -text
|
| 3 |
+
*.bin filter=lfs diff=lfs merge=lfs -text
|
| 4 |
+
*.bz2 filter=lfs diff=lfs merge=lfs -text
|
| 5 |
+
*.ckpt filter=lfs diff=lfs merge=lfs -text
|
| 6 |
+
*.ftz filter=lfs diff=lfs merge=lfs -text
|
| 7 |
+
*.gz filter=lfs diff=lfs merge=lfs -text
|
| 8 |
+
*.h5 filter=lfs diff=lfs merge=lfs -text
|
| 9 |
+
*.joblib filter=lfs diff=lfs merge=lfs -text
|
| 10 |
+
*.lfs.* filter=lfs diff=lfs merge=lfs -text
|
| 11 |
+
*.mlmodel filter=lfs diff=lfs merge=lfs -text
|
| 12 |
+
*.model filter=lfs diff=lfs merge=lfs -text
|
| 13 |
+
*.msgpack filter=lfs diff=lfs merge=lfs -text
|
| 14 |
+
*.npy filter=lfs diff=lfs merge=lfs -text
|
| 15 |
+
*.npz filter=lfs diff=lfs merge=lfs -text
|
| 16 |
+
*.onnx filter=lfs diff=lfs merge=lfs -text
|
| 17 |
+
*.ot filter=lfs diff=lfs merge=lfs -text
|
| 18 |
+
*.parquet filter=lfs diff=lfs merge=lfs -text
|
| 19 |
+
*.pb filter=lfs diff=lfs merge=lfs -text
|
| 20 |
+
*.pickle filter=lfs diff=lfs merge=lfs -text
|
| 21 |
+
*.pkl filter=lfs diff=lfs merge=lfs -text
|
| 22 |
+
*.pt filter=lfs diff=lfs merge=lfs -text
|
| 23 |
+
*.pth filter=lfs diff=lfs merge=lfs -text
|
| 24 |
+
*.rar filter=lfs diff=lfs merge=lfs -text
|
| 25 |
+
*.safetensors filter=lfs diff=lfs merge=lfs -text
|
| 26 |
+
saved_model/**/* filter=lfs diff=lfs merge=lfs -text
|
| 27 |
+
*.tar.* filter=lfs diff=lfs merge=lfs -text
|
| 28 |
+
*.tar filter=lfs diff=lfs merge=lfs -text
|
| 29 |
+
*.tflite filter=lfs diff=lfs merge=lfs -text
|
| 30 |
+
*.tgz filter=lfs diff=lfs merge=lfs -text
|
| 31 |
+
*.wasm filter=lfs diff=lfs merge=lfs -text
|
| 32 |
+
*.xz filter=lfs diff=lfs merge=lfs -text
|
| 33 |
+
*.zip filter=lfs diff=lfs merge=lfs -text
|
| 34 |
+
*.zst filter=lfs diff=lfs merge=lfs -text
|
| 35 |
+
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
| 36 |
+
|
| 37 |
+
# Frozen IEEE notebook: always check out with LF so the SHA-256 freeze check
|
| 38 |
+
# (make freeze-paper-notebook) matches .paper-notebook.sha256 on every OS.
|
| 39 |
+
notebooks/01_ieee_inceptionv3_transformer.ipynb text eol=lf
|
Dockerfile
ADDED
|
@@ -0,0 +1,58 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# syntax=docker/dockerfile:1.7
|
| 2 |
+
# =============================================================================
|
| 3 |
+
# Dockerfile — FastAPI inference backend for HuggingFace Spaces (Docker SDK).
|
| 4 |
+
# -----------------------------------------------------------------------------
|
| 5 |
+
# Target: HF Spaces, hardware = cpu-basic (2 vCPU / 16 GB RAM).
|
| 6 |
+
# Port: 7860 (HF Spaces convention).
|
| 7 |
+
# User: UID 1000 named "user" (HF Spaces requirement).
|
| 8 |
+
# Workdir: /home/user/app (HF Spaces convention).
|
| 9 |
+
# Worker: uvicorn single worker — keeps the TF model loaded once in RAM.
|
| 10 |
+
# =============================================================================
|
| 11 |
+
|
| 12 |
+
FROM python:3.11-slim-bookworm
|
| 13 |
+
|
| 14 |
+
ENV PYTHONDONTWRITEBYTECODE=1 \
|
| 15 |
+
PYTHONUNBUFFERED=1 \
|
| 16 |
+
PIP_DISABLE_PIP_VERSION_CHECK=1 \
|
| 17 |
+
PIP_NO_CACHE_DIR=1 \
|
| 18 |
+
TF_CPP_MIN_LOG_LEVEL=2 \
|
| 19 |
+
HF_HOME=/home/user/.cache/huggingface
|
| 20 |
+
|
| 21 |
+
# libgomp1 is required by tensorflow-cpu (OpenMP runtime).
|
| 22 |
+
# curl is used by HEALTHCHECK.
|
| 23 |
+
RUN apt-get update \
|
| 24 |
+
&& apt-get install -y --no-install-recommends libgomp1 curl \
|
| 25 |
+
&& rm -rf /var/lib/apt/lists/*
|
| 26 |
+
|
| 27 |
+
# HF Spaces requires a non-root user with UID 1000 named "user".
|
| 28 |
+
RUN useradd --create-home --uid 1000 user
|
| 29 |
+
USER user
|
| 30 |
+
ENV PATH="/home/user/.local/bin:${PATH}"
|
| 31 |
+
WORKDIR /home/user/app
|
| 32 |
+
|
| 33 |
+
# --- Dependency layer (cached across code changes) ---------------------------
|
| 34 |
+
COPY --chown=user:user requirements.txt ./
|
| 35 |
+
RUN pip install --user --no-cache-dir -r requirements.txt
|
| 36 |
+
|
| 37 |
+
# --- Application source ------------------------------------------------------
|
| 38 |
+
# Copy only what the runtime needs. Build context is pruned by .dockerignore.
|
| 39 |
+
COPY --chown=user:user pyproject.toml README.md ./
|
| 40 |
+
COPY --chown=user:user src/ ./src/
|
| 41 |
+
COPY --chown=user:user backend/ ./backend/
|
| 42 |
+
COPY --chown=user:user configs/ ./configs/
|
| 43 |
+
COPY --chown=user:user models/ ./models/
|
| 44 |
+
|
| 45 |
+
# Install the local captioning package without re-resolving deps.
|
| 46 |
+
RUN pip install --user --no-cache-dir --no-deps -e .
|
| 47 |
+
|
| 48 |
+
EXPOSE 7860
|
| 49 |
+
|
| 50 |
+
HEALTHCHECK --interval=30s --timeout=10s --start-period=90s --retries=3 \
|
| 51 |
+
CMD curl --fail --silent http://127.0.0.1:7860/healthz || exit 1
|
| 52 |
+
|
| 53 |
+
CMD ["uvicorn", "app.main:app", \
|
| 54 |
+
"--app-dir", "backend", \
|
| 55 |
+
"--host", "0.0.0.0", \
|
| 56 |
+
"--port", "7860", \
|
| 57 |
+
"--workers", "1", \
|
| 58 |
+
"--log-level", "info"]
|
README.md
ADDED
|
@@ -0,0 +1,391 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
title: Image Captioning API
|
| 3 |
+
emoji: 🖼️
|
| 4 |
+
colorFrom: blue
|
| 5 |
+
colorTo: indigo
|
| 6 |
+
sdk: docker
|
| 7 |
+
app_port: 7860
|
| 8 |
+
pinned: false
|
| 9 |
+
license: mit
|
| 10 |
+
short_description: InceptionV3 + Transformer image captioning inference API
|
| 11 |
+
---
|
| 12 |
+
|
| 13 |
+
<h1 align="center">Caption Studio</h1>
|
| 14 |
+
|
| 15 |
+
<p align="center">
|
| 16 |
+
<strong>Upload a photo and get an English caption. The InceptionV3 + Transformer model from an IEEE paper, rebuilt from a research notebook into a tested, deployed full-stack system.</strong>
|
| 17 |
+
</p>
|
| 18 |
+
|
| 19 |
+
<p align="center">
|
| 20 |
+
<a href="https://github.com/apoorvrajdev/image-captioning-system/actions/workflows/ci.yml"><img alt="CI" src="https://github.com/apoorvrajdev/image-captioning-system/actions/workflows/ci.yml/badge.svg?branch=main"></a>
|
| 21 |
+
<img alt="Python 3.10–3.12" src="https://img.shields.io/badge/python-3.10%E2%80%933.12-3776AB?style=flat-square&logo=python&logoColor=white">
|
| 22 |
+
<img alt="TensorFlow 2.15" src="https://img.shields.io/badge/TensorFlow-2.15-FF6F00?style=flat-square&logo=tensorflow&logoColor=white">
|
| 23 |
+
<img alt="FastAPI" src="https://img.shields.io/badge/FastAPI-0.133-009688?style=flat-square&logo=fastapi&logoColor=white">
|
| 24 |
+
<img alt="React 19" src="https://img.shields.io/badge/React-19-61DAFB?style=flat-square&logo=react&logoColor=black">
|
| 25 |
+
<img alt="563 pytest + 27 E2E tests" src="https://img.shields.io/badge/tests-563%20pytest%20%2B%2027%20E2E-success?style=flat-square">
|
| 26 |
+
<a href="https://doi.org/10.1109/ICSPCRE62303.2024.10675203"><img alt="IEEE ICSPCRE 2024" src="https://img.shields.io/badge/IEEE-ICSPCRE%202024-00629B?style=flat-square&logo=ieee&logoColor=white"></a>
|
| 27 |
+
<a href="LICENSE"><img alt="License: MIT" src="https://img.shields.io/badge/license-MIT-blue?style=flat-square"></a>
|
| 28 |
+
</p>
|
| 29 |
+
|
| 30 |
+
<p align="center">
|
| 31 |
+
<a href="https://image-captioning-system.vercel.app"><strong>Live app</strong></a> ·
|
| 32 |
+
<a href="#-product-demo">Demo video</a> ·
|
| 33 |
+
<a href="https://apoorvrajdev-image-captioning-api.hf.space/docs">API docs</a> ·
|
| 34 |
+
<a href="#architecture">Architecture</a> ·
|
| 35 |
+
<a href="docs/EVAL_METHODOLOGY.md">Evaluation methodology</a> ·
|
| 36 |
+
<a href="https://doi.org/10.1109/ICSPCRE62303.2024.10675203">Paper</a>
|
| 37 |
+
</p>
|
| 38 |
+
|
| 39 |
+
## 🎬 Product Demo
|
| 40 |
+
|
| 41 |
+
A short technical demo of the Image Captioning System's live image captioning, model inference pipeline and research-to-production checks.
|
| 42 |
+
|
| 43 |
+
|
| 44 |
+
https://github.com/user-attachments/assets/7d63c76a-019f-4616-a219-fa4043f0233c
|
| 45 |
+
|
| 46 |
+
|
| 47 |
+
**▶ [Watch the 27-second demo]** · [Try the live app](https://image-captioning-system.vercel.app) · [How the demo was made](docs/demo/README.md)
|
| 48 |
+
|
| 49 |
+
**No photo to hand?** Under **Try an example**, the live app offers five licensed COCO images (CC BY 2.0, credited in the app) that go through the same validation and API call as an upload. Pick **Greedy** or **Beam search** before captioning, or press **Compare both** to run the two decoders on the same image one after the other, each with its measured latency. A **Dataset reference** card shows the caption a person wrote for that image in COCO.
|
| 50 |
+
|
| 51 |
+
---
|
| 52 |
+
|
| 53 |
+
Caption Studio describes an image in one sentence. A frozen InceptionV3 network turns the photo into a grid of visual features, and a Transformer decoder writes the caption word by word, either greedily or with beam search. The architecture comes from the IEEE paper *AI Narratives: Bridging Visual Content and Linguistic Expression* and its research notebook, which this repository keeps frozen. Around it sits what it takes for other people to use and maintain that model: an installable Python package, a FastAPI service on Hugging Face Spaces, a React app on Vercel, and CI that gates every deploy.
|
| 54 |
+
|
| 55 |
+
The part worth reading first is the evaluation. With beam search, the deployed checkpoint scored BLEU-4 10.4 against the paper's ~24, and the obvious move was to retrain. A pre-registered, blinded audit showed the gap came mostly from how many reference captions each image is scored against: rescored against all five COCO references, the same captions reach 25.9. The audit also found the real weakness, which is that captions are fluent but often generic. Both findings are [reported below](#what-the-evaluation-taught-me), next to a comparison with three modern captioning models that is caveated rather than ranked.
|
| 56 |
+
|
| 57 |
+
## At a glance
|
| 58 |
+
|
| 59 |
+
| | |
|
| 60 |
+
|---|---|
|
| 61 |
+
| **Product** | React 19 app on Vercel → FastAPI on a Hugging Face Space (Docker, CPU) → one TensorFlow model loaded at startup. Upload a photo or pick an example; choose greedy or beam search, or compare both |
|
| 62 |
+
| **Model** | Frozen InceptionV3 (ImageNet) → 1-layer Transformer encoder → 1-layer, 8-head Transformer decoder, 15,000-token vocabulary, trained on a 120,000-caption COCO 2017 sample. Served as `v2.0.0` from a commit-pinned, SHA-256-verified Hugging Face Hub revision |
|
| 63 |
+
| **Caption quality** | CIDEr **0.789** greedy / **0.826** beam on 500 held-out COCO images with ~1.46 references each. BLEU-4 **25.9** when the beam captions are rescored against all 5 COCO references. A blinded review found 3 of 30 captions image-specific and correct |
|
| 64 |
+
| **Live latency** | About 0.8���1.0 s (greedy) and 1.4–1.5 s (beam) of model inference per caption on the Space's CPU, about 3 s end to end including the upload. A handful of live requests, not a load test |
|
| 65 |
+
| **Software checks** | 563 pytest tests and 27 Playwright browser tests, plus ruff, mypy, a 4-stage notebook parity audit, a SHA-256 notebook freeze, and dependency and secret scans. All of them block CI |
|
| 66 |
+
| **Delivery** | A green push to `main` rebuilds the Space only if the image changed. A deploy passes only after the health check and one real caption from the live model |
|
| 67 |
+
|
| 68 |
+
**Contents:** [What I built](#what-i-built) · [Architecture](#architecture) · [Evaluation](#evaluation) · [Engineering highlights](#engineering-highlights) · [Limitations](#limitations) · [Tech stack](#tech-stack) · [Run locally](#run-locally) · [Testing and CI](#testing-and-ci) · [Repository layout](#repository-layout) · [Roadmap](#roadmap) · [Documentation](#documentation) · [Publication](#publication-and-acknowledgements)
|
| 69 |
+
|
| 70 |
+
---
|
| 71 |
+
|
| 72 |
+
## What I built
|
| 73 |
+
|
| 74 |
+
The paper defined the architecture and its notebook trained the original model. The rest was built in this repository:
|
| 75 |
+
|
| 76 |
+
- **A web app** ([`frontend/`](frontend/)). React 19, Vite 8 and Tailwind v4: drag-and-drop upload with client-side type and size checks, the example gallery, the Greedy / Beam search / Compare both controls, a live API health badge, and a static **Phase 3 comparison** view of the baseline results.
|
| 77 |
+
- **An inference API** ([`backend/app/`](backend/app/)). FastAPI with a Pydantic schema for every request and response, one warm model per process, TensorFlow work moved off the event loop, bounded admission, and a status-code contract shared with the frontend.
|
| 78 |
+
- **A modular ML package** ([`src/captioning/`](src/captioning/)). The notebook split into config, preprocessing, data, models, training, greedy and beam-search inference, and evaluation, with a parity audit that proves the split didn't change behaviour.
|
| 79 |
+
- **A retrained checkpoint.** The stabilized recipe (label smoothing, warmup into cosine decay), trained on COCO 2017 on Kaggle and published to the Hub as `v2.0.0`. This is what the live app serves.
|
| 80 |
+
- **Evaluation tooling.** BLEU-1–4, METEOR, ROUGE-L and CIDEr on a fixed slice, one machine-readable artefact set per run, and a harness that compares the model with BLIP, ViT-GPT2 and GIT under a protocol committed before any result.
|
| 81 |
+
- **Delivery.** A Docker image on Hugging Face Spaces, Vercel for the frontend, and GitHub Actions for CI, security scanning and gated deploys.
|
| 82 |
+
|
| 83 |
+
## Architecture
|
| 84 |
+
|
| 85 |
+
```mermaid
|
| 86 |
+
flowchart TB
|
| 87 |
+
subgraph live["Serving path (live)"]
|
| 88 |
+
SPA["React 19 SPA · Vercel<br/>upload · example gallery<br/>greedy / beam / compare both"]
|
| 89 |
+
API["FastAPI · Hugging Face Space<br/>Docker, CPU, 1 worker<br/>GET /healthz · POST /v1/captions"]
|
| 90 |
+
PRED["CaptionPredictor<br/>one warm model · 2 inference slots"]
|
| 91 |
+
SPA -- "multipart image + decode_strategy" --> API --> PRED
|
| 92 |
+
end
|
| 93 |
+
|
| 94 |
+
subgraph offline["Research and evaluation"]
|
| 95 |
+
NB["IEEE notebook<br/><i>frozen, SHA-256 locked</i>"]
|
| 96 |
+
PKG["src/captioning<br/>config · data · models · training<br/>inference · evaluation"]
|
| 97 |
+
TRAIN["Training on Kaggle<br/>configs/train/stabilized.yaml"]
|
| 98 |
+
RES[("results/<run_id>/<br/>metrics · predictions · audits")]
|
| 99 |
+
NB -- "4-stage parity audit in CI" --- PKG
|
| 100 |
+
PKG --> TRAIN
|
| 101 |
+
PKG -- "evaluate · compare · benchmark" --> RES
|
| 102 |
+
end
|
| 103 |
+
|
| 104 |
+
TRAIN -- "publish v2.0.0" --> HUB[("Hugging Face Hub<br/>model.h5 · vocab.json")]
|
| 105 |
+
HUB -. "pinned commit, SHA-256 checked at startup" .-> PRED
|
| 106 |
+
RES -. "exported to JSON at build time" .-> SPA
|
| 107 |
+
```
|
| 108 |
+
|
| 109 |
+
The two halves meet in two places: the Hub weights the API downloads at startup, and the results JSON the SPA bundles at build time. The serving image contains no notebook, no training data and no evaluation results. The notebook is the published research artefact and stays byte-for-byte as it was; every improvement goes into the package, behind the parity audit.
|
| 110 |
+
|
| 111 |
+
### Inference pipeline
|
| 112 |
+
|
| 113 |
+
```text
|
| 114 |
+
POST /v1/captions multipart: image, optional decode_strategy
|
| 115 |
+
→ body cap 413 before a body over 10 MiB (+64 KiB framing) is buffered
|
| 116 |
+
→ decode_strategy greedy (default) or beam, otherwise 422
|
| 117 |
+
→ content-type allow-list JPEG · PNG · BMP, otherwise 415 (WebP and GIF included)
|
| 118 |
+
→ bounded read 400 if empty, 413 if over 10 MiB
|
| 119 |
+
→ admission 503 + Retry-After: 2 while both inference slots are busy
|
| 120 |
+
→ header check + decode 413 over 16,384 px a side or 40 MP, before TensorFlow decodes; 422 if undecodable
|
| 121 |
+
→ preprocess_image_tensor 299×299, InceptionV3 scaling: the same function training uses
|
| 122 |
+
→ InceptionV3 (frozen) 8×8×2048 feature map → 64 image tokens
|
| 123 |
+
→ Transformer encoder 1 layer, 1 head, 512-dim
|
| 124 |
+
→ Transformer decoder 1 layer, 8 heads; greedy, or beam search (width 4)
|
| 125 |
+
→ CaptionResponse caption · model_version · decode_strategy · latency_ms · request_id
|
| 126 |
+
```
|
| 127 |
+
|
| 128 |
+
Beam search runs with exactly the settings of the evaluated beam run: width 4, length penalty 0.7, repetition penalty 1.2, no repeated trigrams. A live beam caption is therefore the decoder the metrics below describe.
|
| 129 |
+
|
| 130 |
+
| Endpoint | Behaviour |
|
| 131 |
+
|---|---|
|
| 132 |
+
| [`GET /healthz`](https://apoorvrajdev-image-captioning-api.hf.space/healthz) | Always 200; `model_loaded` and `model_version` report readiness. Runs on the event loop, so it never waits behind inference |
|
| 133 |
+
| `POST /v1/captions` | Image in, `CaptionResponse` out. Status codes 200 / 400 / 413 / 415 / 422 / 503, each covered by a backend test and consumed by [`frontend/src/services/api.js`](frontend/src/services/api.js) |
|
| 134 |
+
| [`GET /docs`](https://apoorvrajdev-image-captioning-api.hf.space/docs) | Swagger UI generated from the Pydantic schemas |
|
| 135 |
+
|
| 136 |
+
**Where it runs.** The SPA is on [Vercel](https://image-captioning-system.vercel.app), rebuilt on every push to `main`. The API is a [Hugging Face Space](https://huggingface.co/spaces/apoorvrajdev/image-captioning-api): Docker, `python:3.11-slim`, a non-root user, one uvicorn worker on a CPU instance. The weights are in the Hub repo [`apoorvrajdev/captioning-inceptionv3-transformer`](https://huggingface.co/apoorvrajdev/captioning-inceptionv3-transformer) at tag `v2.0.0`, which the backend resolves to commit `59d93b4`. Topology, promotion and rollback are in the [deployment runbook](docs/PHASE_2C_DEPLOYMENT_RUNBOOK.md).
|
| 137 |
+
|
| 138 |
+
## Evaluation
|
| 139 |
+
|
| 140 |
+
Every number in this section comes from a committed run directory under [`results/`](results/). The evaluation slice is 500 COCO 2017 `train2017` images from the model's own image-level validation split, so the model never trained on them. They are scored against the 732 reference captions stored with them, about 1.46 per image. That is fewer than the standard five per image, which lowers BLEU in particular; [the audit below](#what-the-evaluation-taught-me) measures by how much.
|
| 141 |
+
|
| 142 |
+
### Caption quality
|
| 143 |
+
|
| 144 |
+
| Decoding | BLEU-1 | BLEU-4 | METEOR | ROUGE-L | CIDEr |
|
| 145 |
+
|---|---|---|---|---|---|
|
| 146 |
+
| Greedy ([run](results/stabilized-greedy/)) | 42.20 | 10.57 | 15.45 | 37.57 | 0.789 |
|
| 147 |
+
| Beam search, width 4 ([run](results/stabilized-beam-w4-lp07-rp12/)) | 41.93 | 10.39 | 15.56 | 36.84 | **0.826** |
|
| 148 |
+
|
| 149 |
+
Beam search gives slightly less n-gram overlap but about 5% more CIDEr, the metric that rewards image-specific wording over generic phrases. That is why the app lets you choose. BLEU-2 and BLEU-3, predictions and per-sample diagnostics are in each run directory.
|
| 150 |
+
|
| 151 |
+
The served checkpoint used the stabilized recipe ([`configs/train/stabilized.yaml`](configs/train/stabilized.yaml)): the notebook's architecture and hyperparameters, plus label smoothing 0.1, a 500-step warmup into cosine decay, and dropout off during validation. It trained on 120,000 COCO 2017 captions, split 80/20 by image, for 10 epochs at batch 64. Training loss fell from 4.69 to 3.33, and validation accuracy rose every epoch, from 0.429 to 0.475 ([`history.json`](results/history.json)).
|
| 152 |
+
|
| 153 |
+
### Sample captions
|
| 154 |
+
|
| 155 |
+
These are the five example images in the live app, with the served `v2.0.0` checkpoint's captions from the committed evaluation runs. The kitten captions are the ones the deployed app returned in the [demo recording](docs/demo/README.md).
|
| 156 |
+
|
| 157 |
+
| Image | Greedy | Beam search | COCO reference (written by a person) |
|
| 158 |
+
|---|---|---|---|
|
| 159 |
+
| <img src="https://raw.githubusercontent.com/apoorvrajdev/image-captioning-system/main/frontend/public/samples/000000418799.jpg" width="140" alt="Kitten on a computer keyboard"> | a cat is sitting on a keyboard of a computer | a close up of a cat laying on a keyboard | A kitten sitting on the keyboard of a computer. |
|
| 160 |
+
| <img src="https://raw.githubusercontent.com/apoorvrajdev/image-captioning-system/main/frontend/public/samples/000000316059.jpg" width="140" alt="Red fire hydrant on a grass verge"> | a red fire hydrant sitting on a sidewalk | a red fire hydrant sitting on the side of a road | A red fire hydrant sitting on the side of a road. |
|
| 161 |
+
| <img src="https://raw.githubusercontent.com/apoorvrajdev/image-captioning-system/main/frontend/public/samples/000000215132.jpg" width="140" alt="Living room with a sofa and fireplace"> | a living room with a couch and a tv | a living room filled with furniture and a flat screen tv | a small living room with a blue couch, flat screen television and a fireplace |
|
| 162 |
+
| <img src="https://raw.githubusercontent.com/apoorvrajdev/image-captioning-system/main/frontend/public/samples/000000422807.jpg" width="140" alt="Propeller aircraft in a blue sky"> | a large airplane flying through the sky | an airplane is flying in the sky on a cloudy day | An airplane flying through a blue sky with propellers. |
|
| 163 |
+
| <img src="https://raw.githubusercontent.com/apoorvrajdev/image-captioning-system/main/frontend/public/samples/000000302885.jpg" width="140" alt="Platter of sliced fruit"> | a plate of fruit and a salad with a fork | a close up of a plate of food on a table | A collection of fruits are stacked on a plate. |
|
| 164 |
+
|
| 165 |
+
Neither decoder wins every image. Beam search matches the reference wording on the hydrant and adds the flat-screen TV in the living room, but it calls a clear blue sky cloudy and is vaguer than greedy on the fruit platter.
|
| 166 |
+
|
| 167 |
+
<sub>Images: COCO 2017 via Flickr, <a href="https://creativecommons.org/licenses/by/2.0/">CC BY 2.0</a>: <a href="https://www.flickr.com/photos/slava/496607907/">“Mango HaX0ring”</a> by Slava · <a href="https://www.flickr.com/photos/pa1nt/2707331770/">“07252008”</a> by pa1nt · <a href="https://www.flickr.com/photos/amy_elizabeth_west/5067711955/">“Living room w/painted cable covers”</a> by Amy West · <a href="https://www.flickr.com/photos/39551170@N02/6815669873/">“US Coast Guard Herc”</a> by Simon_sees · <a href="https://www.flickr.com/photos/little_yiye/3742500176/">“fruit platter”</a> by >littleyiye<. Reference captions: COCO Consortium, <a href="https://creativecommons.org/licenses/by/4.0/">CC BY 4.0</a>.</sub>
|
| 168 |
+
|
| 169 |
+
### What the evaluation taught me
|
| 170 |
+
|
| 171 |
+
**Before retraining a model to close a gap, check whether the gap is real.**
|
| 172 |
+
|
| 173 |
+
The beam checkpoint's BLEU-4 of 10.39 looked about 14 points behind the ~24 the paper reports. Instead of spending GPU hours on a retrain, I ran a pre-registered, blinded audit ([`EVAL_METHODOLOGY.md`](docs/EVAL_METHODOLOGY.md) §§ 1–7). The hypotheses and decision thresholds were committed before any result, and the qualitative review was judged before the BLEU result was unblinded.
|
| 174 |
+
|
| 175 |
+
| The same 500 beam captions, scored against | References per image | BLEU-4 |
|
| 176 |
+
|---|---|---|
|
| 177 |
+
| The references stored with the slice | ~1.46 | 10.39 |
|
| 178 |
+
| All COCO references for those images | 5 | **25.91** |
|
| 179 |
+
|
| 180 |
+
BLEU credits an n-gram that matches *any* reference, so the reference count alone moved the headline number by 15 points. Scored on the paper's footing, the checkpoint lands in the same range as the paper. That is **methodology parity, not a claim of beating the paper**. The audit also measured a smaller effect from BLEU smoothing and aggregation: NLTK's variants score the same captions 22.2–23.8 ([`metrics_5ref.json`](results/stabilized-beam-w4-lp07-rp12/metrics_5ref.json)).
|
| 181 |
+
|
| 182 |
+
The blinded review of 30 beam captions found the real weakness:
|
| 183 |
+
|
| 184 |
+
| Category | Captions |
|
| 185 |
+
|---|---|
|
| 186 |
+
| Specific and correct: the right subject plus a distinguishing detail | 3 |
|
| 187 |
+
| Correct but generic ("a person on a beach") | 11 |
|
| 188 |
+
| Partly correct: a count, colour, action or secondary object wrong | 15 |
|
| 189 |
+
| Incorrect | 1 |
|
| 190 |
+
|
| 191 |
+
Captions are fluent and usually on topic, but mostly generic. With 30 samples, each proportion is uncertain by up to about ±18 percentage points. The verdict was **reframe, do not retrain** ([`verdict.md`](results/stabilized-beam-w4-lp07-rp12/verdict.md)): the BLEU gap was mostly the evaluation setup, and specificity is limited by the frozen ImageNet encoder, which the original training recipe would not change. Retraining on the paper's recipe stays on the [roadmap](#roadmap) as an experiment, not a fix.
|
| 192 |
+
|
| 193 |
+
### Comparison with BLIP, ViT-GPT2 and GIT
|
| 194 |
+
|
| 195 |
+
Phase 3 scored three Hugging Face captioning models on the same 500 images and references, at pinned revisions with greedy decoding. The protocol was committed before the first baseline ran ([`EVAL_METHODOLOGY.md`](docs/EVAL_METHODOLOGY.md) § 8).
|
| 196 |
+
|
| 197 |
+
> **Read this table as context, not a ranking.** The slice is held out for this repository's model only. BLIP's and GIT's model cards state COCO training, and ViT-GPT2 may have used it, so the baselines may have seen these exact images; per-image overlap isn't measured (§ 8.5). With ~1.46 references per image, none of these scores is comparable with published COCO results either.
|
| 198 |
+
|
| 199 |
+
| Model | Decoding | BLEU-4 | METEOR | ROUGE-L | CIDEr |
|
| 200 |
+
|---|---|---|---|---|---|
|
| 201 |
+
| BLIP-base | greedy | 19.88 | 17.23 | 42.13 | 1.06 |
|
| 202 |
+
| GIT-base-coco | greedy | 18.83 | 21.92 | 47.08 | 1.46 |
|
| 203 |
+
| CNN + Transformer (this repo) | beam | 10.39 | 15.56 | 36.84 | 0.83 |
|
| 204 |
+
| CNN + Transformer (this repo) | greedy | 10.57 | 15.45 | 37.57 | 0.79 |
|
| 205 |
+
| ViT-GPT2 | greedy | 15.84 | 19.78 | 44.51 | 1.26 |
|
| 206 |
+
|
| 207 |
+
Rows follow [`comparison.md`](results/phase3-comparison/comparison.md), ordered by model id. All three baselines score higher on this slice. Because they may have trained on these images, the table can't say how large that gap is on genuinely unseen images. What it does establish is that the harness reproduces this repository's own results exactly: re-running the greedy row through it gave 500/500 identical captions and bit-identical metrics (§ 8.8).
|
| 208 |
+
|
| 209 |
+
Latency for all four models, on CPU (a local laptop) and on a Kaggle T4 GPU, is in §§ 9.9–9.10 and in the app's **Phase 3 comparison** view. The two devices were different hosts, so the figures compare models within a device, not CPU against GPU. This model's batch-8 figure is eight sequential calls rather than one batched call.
|
| 210 |
+
|
| 211 |
+
## Engineering highlights
|
| 212 |
+
|
| 213 |
+
- **The notebook is frozen, and parity is tested, not assumed.** The paper's notebook stays byte-for-byte as published, locked by a SHA-256 check in CI. A four-stage audit ([`scripts/notebook_module_audit.py`](scripts/notebook_module_audit.py)) re-implements caption preprocessing, tokenisation, image preprocessing and the decoder forward pass from the notebook, and asserts the package matches them. Every later improvement is opt-in behind a flag whose default reproduces the notebook, so a changed number always traces back to one named change.
|
| 214 |
+
- **Training and serving share one preprocessing path.** Uploads go through the same `preprocess_image_tensor` that builds the training pipeline, so the model never sees pixels normalised differently in production than in training.
|
| 215 |
+
- **Uploads are bounded before they cost anything.** A middleware refuses oversized bodies before the multipart parser buffers them, and the route reads at most one byte past the limit. A header check rejects images over 16,384 px a side or 40 megapixels before TensorFlow decodes them. WebP is refused at the door because TensorFlow 2.15 can't decode it, instead of failing deep inside the model.
|
| 216 |
+
- **Inference is bounded too.** The Space runs one CPU worker. At most two captions decode at once, and the next request gets an immediate 503 with `Retry-After: 2` rather than joining a queue that slows everyone down. `/healthz` runs on the event loop, so a busy model never fails the health check.
|
| 217 |
+
- **The weights are pinned and verified.** Hub tags can be moved, so the backend resolves `v2.0.0` to a commit id written in reviewed code. It downloads only `model.h5` and `vocab.json` and checks both against SHA-256 digests pinned per commit before anything reads them. The vocabulary loads from validated JSON instead of unpickling a file from the network ([ADR-032, ADR-033](docs/DECISIONS.md)).
|
| 218 |
+
- **Decoding is a per-request choice on one model.** Greedy is the default, and beam search uses the evaluated settings. Compare both in the app is two ordinary requests to the same loaded weights, not a second model.
|
| 219 |
+
- **Configuration is strict.** Every config model is Pydantic v2 with `extra="forbid"`, so a misspelt hyperparameter fails at load instead of silently falling back to a default. Research settings (`CAPTIONING__*`) and deployment settings (`BACKEND_*`) are separate objects, and environment variables outrank the YAML ([ADR-029](docs/DECISIONS.md)).
|
| 220 |
+
- **Only tested commits deploy, and only when needed.** After green CI, the Space rebuilds only if an image input changed since the last successful deploy, and it receives a single-commit snapshot of exactly the tested build context. A deploy passes only once the health check reports the model loaded and one real caption comes back from the live Space ([ADR-027, ADR-028, ADR-030](docs/DECISIONS.md)).
|
| 221 |
+
- **Results are artefacts, not prose.** Every evaluation writes a run directory with metrics, predictions and per-sample diagnostics, and run directories are never rewritten. The app's comparison view is exported from those files at build time, and a test fails if the two disagree ([ADR-021](docs/DECISIONS.md)).
|
| 222 |
+
- **TensorFlow is pinned at 2.15 on purpose.** TensorFlow 2.16 switches to Keras 3, which breaks `TextVectorization` save and load, and serving depends on that tokenizer round trip. The baselines' `torch` and `transformers` live in an optional `[hf]` extra and are imported lazily, so the serving image never installs them.
|
| 223 |
+
|
| 224 |
+
## Limitations
|
| 225 |
+
|
| 226 |
+
### Model
|
| 227 |
+
|
| 228 |
+
- Captions are often generic, and count, colour and attribute mistakes are common (see the blinded review above, and [`qualitative.jsonl`](results/stabilized-beam-w4-lp07-rp12/qualitative.jsonl) for 30 predictions next to their references). This is a frozen ImageNet CNN with a one-layer decoder, trained for 10 epochs on a 120,000-caption sample: a working demonstration of the paper's model, not a general-purpose captioner.
|
| 229 |
+
- It writes English only, in the vocabulary of COCO's everyday scenes. Unusual compositions fall back to generic captions.
|
| 230 |
+
- Quality is measured on one 500-image slice with ~1.46 references per image. The 5-reference rescore covers BLEU only.
|
| 231 |
+
- The baseline comparison is not held out for the baselines, as explained above.
|
| 232 |
+
|
| 233 |
+
### System
|
| 234 |
+
|
| 235 |
+
- It is a public demo with no authentication and no rate limiting. One heavy client can keep both inference slots busy, and everyone else gets the 503 until a slot frees.
|
| 236 |
+
- It runs on a CPU and one worker. The first request after an idle period can be slow while the Space wakes up.
|
| 237 |
+
- The Hugging Face proxy answers CORS itself and reflects any origin, so the app's allowed-origin list isn't enforced at the edge.
|
| 238 |
+
- The frontend has no unit tests. The browser tests run on Chromium only, against a mocked API.
|
| 239 |
+
|
| 240 |
+
The full list, with risks and suggested fixes, is in [`SECURITY.md`](docs/SECURITY.md) § Known gaps and [`PHASE_1_NOTES.md`](docs/PHASE_1_NOTES.md#technical-debt-remaining).
|
| 241 |
+
|
| 242 |
+
## Tech stack
|
| 243 |
+
|
| 244 |
+
| Area | Technologies |
|
| 245 |
+
|---|---|
|
| 246 |
+
| **Frontend** | React 19, Vite 8, Tailwind CSS v4, ESLint (flat config) |
|
| 247 |
+
| **Backend** | FastAPI 0.133, Starlette 1.3, Pydantic v2 and pydantic-settings, Uvicorn, structlog, anyio |
|
| 248 |
+
| **ML and inference** | TensorFlow-CPU 2.15 (pinned), NumPy 1.26, Pillow, `huggingface-hub` |
|
| 249 |
+
| **Evaluation** | sacrebleu (BLEU), pycocoevalcap (CIDEr, METEOR), rouge-score (ROUGE-L); optional `[hf]` extra with `torch` 2.3.0 and `transformers` 4.41.2 for the baselines |
|
| 250 |
+
| **Testing and quality** | pytest, Playwright (Chromium), ruff, mypy, pre-commit (including nbstripout and gitleaks), pip-audit, npm audit |
|
| 251 |
+
| **Deployment** | Docker, Hugging Face Spaces and Hub, Vercel, GitHub Actions |
|
| 252 |
+
|
| 253 |
+
## Run locally
|
| 254 |
+
|
| 255 |
+
**Prerequisites:** Python 3.10–3.12 (TensorFlow 2.15 has no 3.13 wheels), Node 20.19+ or 22.12+ (the range Vite 8 supports), and Git.
|
| 256 |
+
|
| 257 |
+
```bash
|
| 258 |
+
# Backend: install (on Windows, activate with .venv\Scripts\activate)
|
| 259 |
+
python3.10 -m venv .venv && source .venv/bin/activate
|
| 260 |
+
pip install -r requirements-dev.txt -r requirements-eval.txt
|
| 261 |
+
pip install -e ".[hf,mlflow]"
|
| 262 |
+
|
| 263 |
+
# Serve the published v2.0.0 weights: downloaded from the Hub on first start, SHA-256-checked before loading
|
| 264 |
+
export BACKEND_WEIGHTS_HUB_REPO=apoorvrajdev/captioning-inceptionv3-transformer
|
| 265 |
+
export BACKEND_WEIGHTS_HUB_REVISION=v2.0.0
|
| 266 |
+
export BACKEND_MODEL_VERSION=v2.0.0
|
| 267 |
+
uvicorn app.main:app --app-dir backend --port 8000 # Swagger UI at http://localhost:8000/docs
|
| 268 |
+
|
| 269 |
+
# Frontend, in a second terminal
|
| 270 |
+
cd frontend
|
| 271 |
+
npm install
|
| 272 |
+
npm run dev # http://localhost:5173
|
| 273 |
+
```
|
| 274 |
+
|
| 275 |
+
The SPA calls `http://127.0.0.1:8000` unless `VITE_API_BASE` is set (see [`frontend/.env.example`](frontend/.env.example)). To caption from the command line:
|
| 276 |
+
|
| 277 |
+
```bash
|
| 278 |
+
curl -F "image=@photo.jpg" -F "decode_strategy=beam" http://localhost:8000/v1/captions
|
| 279 |
+
```
|
| 280 |
+
|
| 281 |
+
Without network access, [`scripts/bootstrap_dev_artifacts.py`](scripts/bootstrap_dev_artifacts.py) writes untrained weights that boot the whole stack; their captions are gibberish by design. Training needs COCO 2017 locally: `python -m scripts.train --config configs/train/stabilized.yaml --output-dir outputs/runs/stabilized`, with the full procedure in the [stabilized training runbook](docs/STABILIZED_TRAINING_RUNBOOK.md). Deployment, weights promotion and rollback are in the [deployment runbook](docs/PHASE_2C_DEPLOYMENT_RUNBOOK.md).
|
| 282 |
+
|
| 283 |
+
## Testing and CI
|
| 284 |
+
|
| 285 |
+
| Suite | Command | Tests |
|
| 286 |
+
|---|---|---|
|
| 287 |
+
| Library and script unit tests | `pytest tests -q` | 446 |
|
| 288 |
+
| Backend route tests (fake predictor, no TensorFlow) | `pytest backend/app/tests -q` | 117 |
|
| 289 |
+
| Browser E2E (Playwright, Chromium, API mocked) | `cd frontend && npm run test:e2e` | 27 |
|
| 290 |
+
| Notebook parity audit | `python -m scripts.notebook_module_audit` | 4 stages |
|
| 291 |
+
|
| 292 |
+
Every test is CPU-only and offline: no network, no model download. Playwright needs `npx playwright install chromium` once per machine. The exact lint and type-check invocations are in [`docs/CI.md`](docs/CI.md) § Local equivalents.
|
| 293 |
+
|
| 294 |
+
CI ([`ci.yml`](.github/workflows/ci.yml)) runs on every push and pull request to `main`, and every job blocks:
|
| 295 |
+
|
| 296 |
+
| Job | Checks |
|
| 297 |
+
|---|---|
|
| 298 |
+
| Python quality | ruff lint and format check, mypy |
|
| 299 |
+
| pytest (Python 3.10 and 3.11) | all 563 tests, then the parity audit |
|
| 300 |
+
| Notebook freeze | the notebook's SHA-256 against `.paper-notebook.sha256` |
|
| 301 |
+
| pre-commit | every hook on every file |
|
| 302 |
+
| Security | pip-audit of `requirements.txt` against a reviewed baseline, gitleaks over the full history |
|
| 303 |
+
| Frontend | ESLint, production build, Playwright E2E, `npm audit --omit=dev` |
|
| 304 |
+
|
| 305 |
+
After a green run on `main`, [`deploy-backend.yml`](.github/workflows/deploy-backend.yml) decides whether the Space image changed, deploys the tested commit if it did, and gates on the health check and a real caption. Vercel deploys the frontend through its Git integration.
|
| 306 |
+
|
| 307 |
+
## Repository layout
|
| 308 |
+
|
| 309 |
+
```text
|
| 310 |
+
image-captioning-system/
|
| 311 |
+
├── notebooks/ # the paper's notebook, frozen (hash in .paper-notebook.sha256)
|
| 312 |
+
├── src/captioning/ # installable package
|
| 313 |
+
│ ├── config/ # Pydantic v2 schema, YAML + environment loader
|
| 314 |
+
│ ├── preprocessing/ # caption cleanup, image preprocessing shared by training and serving, tokenizer
|
| 315 |
+
│ ├── data/ # COCO loading, image-level split, tf.data pipeline
|
| 316 |
+
│ ├── models/ # InceptionV3 encoder, Transformer encoder and decoder, model factory
|
| 317 |
+
│ ├── training/ # losses, learning-rate schedule, trainer
|
| 318 |
+
│ ├── inference/ # CaptionPredictor, greedy and beam-search decoders
|
| 319 |
+
│ ├── evaluation/ # metrics, run artefacts, Phase 3 comparison and latency
|
| 320 |
+
│ └── baselines/ # one Captioner interface for this model and the Hugging Face baselines
|
| 321 |
+
├── backend/app/ # FastAPI: routes, schemas, services (predictor, weights loader), settings, tests
|
| 322 |
+
├── frontend/ # React SPA; src/services/api.js is the only API boundary, e2e/ the Playwright suite
|
| 323 |
+
├── configs/ # base.yaml reproduces the notebook; train/stabilized.yaml is the served recipe
|
| 324 |
+
├── scripts/ # train, evaluate, predict, compare, benchmark, parity audit, deploy helpers
|
| 325 |
+
├── results/ # committed evaluation runs, append-only
|
| 326 |
+
├── tests/unit/ # library and script tests
|
| 327 |
+
├── docs/ # methodology, runbooks, decisions, security, task history
|
| 328 |
+
├── Dockerfile # the Hugging Face Space image
|
| 329 |
+
└── .github/workflows/ # CI and the gated Space deploy
|
| 330 |
+
```
|
| 331 |
+
|
| 332 |
+
## Roadmap
|
| 333 |
+
|
| 334 |
+
Phases 0–4 are complete: modularisation behind a parity gate, training stabilisation and the metric suite, the API and web app, public deployment, the baseline comparison, and production hardening. Work since then has hardened the upload path and the weights supply chain and added the example gallery and per-request decoding.
|
| 335 |
+
|
| 336 |
+
<details>
|
| 337 |
+
<summary>Completed phases</summary>
|
| 338 |
+
|
| 339 |
+
- **Phase 0, bootstrap.** Packaging, pre-commit hooks, the notebook freeze and pinned dependencies.
|
| 340 |
+
- **Phase 1, modularisation.** The notebook split into the `captioning` package behind the four-stage parity audit, with strict typed configuration and unit tests.
|
| 341 |
+
- **Phase 1b, training stabilisation.** Beam search, the BLEU / METEOR / ROUGE-L / CIDEr suite with per-run artefacts, the opt-in stabilized recipe and the `v2.0.0` checkpoint, then the Stage 0 evaluation audit (reframe, do not retrain).
|
| 342 |
+
- **Phase 2, serving and deployment.** The FastAPI service (2A), the React SPA (2B), and public deployment on Hugging Face Spaces and Vercel with CI/CD (2C).
|
| 343 |
+
- **Phase 3, multimodal baselines.** The pre-registered comparison with BLIP, ViT-GPT2 and GIT, CPU and GPU latency, the static comparison view, and Playwright E2E.
|
| 344 |
+
- **Phase 4, production hardening.** A pinned CI platform, upgraded serving dependencies, dependency and secret scanning, deploys only when the image changes, and the post-deploy caption check.
|
| 345 |
+
- **Since Phase 4.** Environment variables that outrank the YAML, single-commit Space snapshots, pre-decode pixel limits, bounded inference admission, commit-pinned and SHA-256-verified weights, the example gallery, and per-request greedy or beam decoding with Compare both.
|
| 346 |
+
|
| 347 |
+
Task-by-task history, with the verification behind each task: [`docs/TASKS.md`](docs/TASKS.md).
|
| 348 |
+
|
| 349 |
+
</details>
|
| 350 |
+
|
| 351 |
+
**Possible next steps.** These are directions, not current capabilities:
|
| 352 |
+
|
| 353 |
+
- **Original-recipe ablation.** Retrain with the paper's recipe (constant Adam, no label smoothing, early stopping) under the same evaluation, to find out whether it beats the stabilized recipe on this architecture.
|
| 354 |
+
- **A stronger encoder.** Specificity is limited by the frozen InceptionV3. Fine-tuning a modern vision-language captioner such as BLIP-2 on COCO, under the same protocol, would show what better captions cost per request.
|
| 355 |
+
- **Streaming captions.** Server-sent events from `/v1/captions`, so the app shows words as the decoder produces them.
|
| 356 |
+
- **Batch captioning.** One request for a list of images, to amortise per-request overhead.
|
| 357 |
+
- **Visual question answering.** A `POST /v1/vqa` route that reuses the warm encoder.
|
| 358 |
+
- **Online evaluation.** Score each new checkpoint on a held-out slice automatically and publish quality and latency over time.
|
| 359 |
+
|
| 360 |
+
## Documentation
|
| 361 |
+
|
| 362 |
+
| Document | What it covers |
|
| 363 |
+
|---|---|
|
| 364 |
+
| [`docs/EVAL_METHODOLOGY.md`](docs/EVAL_METHODOLOGY.md) | The Stage 0 audit (§§ 1–7), the Phase 3 comparison protocol and results (§ 8), and the latency protocol and results (§ 9) |
|
| 365 |
+
| [`docs/PHASE_2C_DEPLOYMENT_RUNBOOK.md`](docs/PHASE_2C_DEPLOYMENT_RUNBOOK.md) | Deployment topology, Space and Vercel setup, weights promotion and rollback |
|
| 366 |
+
| [`docs/CI.md`](docs/CI.md) | Every CI and deploy job, and their local equivalents |
|
| 367 |
+
| [`docs/SECURITY.md`](docs/SECURITY.md) | Controls in place, known gaps, the dependency audit and the CI scanning policy |
|
| 368 |
+
| [`docs/DECISIONS.md`](docs/DECISIONS.md) | Architecture decision records |
|
| 369 |
+
| [`docs/TEST_PLAN.md`](docs/TEST_PLAN.md) | The checks each kind of change must pass |
|
| 370 |
+
| [`docs/STABILIZED_TRAINING_RUNBOOK.md`](docs/STABILIZED_TRAINING_RUNBOOK.md) | How the stabilized checkpoint was trained |
|
| 371 |
+
| [`docs/TASKS.md`](docs/TASKS.md) · [`docs/MEMORY.md`](docs/MEMORY.md) | Task history with verification records · current project state |
|
| 372 |
+
| [`notebooks/README.md`](notebooks/README.md) | The frozen-notebook policy |
|
| 373 |
+
| [`docs/demo/README.md`](docs/demo/README.md) | How the demo video was made, what in it is real, and its credits |
|
| 374 |
+
| [`docs/PHASE_0_NOTES.md`](docs/PHASE_0_NOTES.md) · [`docs/PHASE_1_NOTES.md`](docs/PHASE_1_NOTES.md) · [`docs/restructure-plan.md`](docs/restructure-plan.md) | Early phase notes and the original restructure plan |
|
| 375 |
+
|
| 376 |
+
## Publication and acknowledgements
|
| 377 |
+
|
| 378 |
+
The model architecture and the original notebook come from:
|
| 379 |
+
|
| 380 |
+
> Preetam, S. C. Muppalla, A. Raj and J. Chawla, "AI Narratives: Bridging Visual Content and Linguistic Expression," *2024 IEEE International Conference on Smart Power Control and Renewable Energy (ICSPCRE)*, 2024. doi: [10.1109/ICSPCRE62303.2024.10675203](https://doi.org/10.1109/ICSPCRE62303.2024.10675203) · [IEEE Xplore](https://ieeexplore.ieee.org/document/10675203)
|
| 381 |
+
|
| 382 |
+
The notebook is kept verbatim at [`notebooks/01_ieee_inceptionv3_transformer.ipynb`](notebooks/01_ieee_inceptionv3_transformer.ipynb). Improvements go into `src/captioning/`, never into the notebook.
|
| 383 |
+
|
| 384 |
+
- **[COCO 2017](https://cocodataset.org/)** provides the training and evaluation data. The annotations are CC BY 4.0, and each image keeps its own Flickr licence.
|
| 385 |
+
- **InceptionV3** ImageNet weights come from Keras Applications.
|
| 386 |
+
- **BLIP-base** (Salesforce), **ViT-GPT2** (nlpconnect) and **GIT-base-coco** (Microsoft), from the Hugging Face Hub, are used for evaluation only and are never served.
|
| 387 |
+
- The demo video's music, sound effects and typefaces are credited in [`docs/demo/README.md`](docs/demo/README.md#credits).
|
| 388 |
+
|
| 389 |
+
## License
|
| 390 |
+
|
| 391 |
+
Released under the [MIT License](LICENSE). Built by [apoorvrajdev](https://github.com/apoorvrajdev) · [apoorvrajmgr@gmail.com](mailto:apoorvrajmgr@gmail.com)
|
backend/app/__init__.py
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
"""FastAPI inference backend for the captioning ML package."""
|
backend/app/api/__init__.py
ADDED
|
@@ -0,0 +1,5 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""HTTP routes for the captioning service."""
|
| 2 |
+
|
| 3 |
+
from app.api.routes import router
|
| 4 |
+
|
| 5 |
+
__all__ = ["router"]
|
backend/app/api/routes.py
ADDED
|
@@ -0,0 +1,170 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""HTTP routes: ``/healthz`` and ``/v1/captions``.
|
| 2 |
+
|
| 3 |
+
Routes are intentionally thin: validate inputs, delegate to the
|
| 4 |
+
``PredictorService``, shape the response. No model code, no TF imports.
|
| 5 |
+
"""
|
| 6 |
+
|
| 7 |
+
from __future__ import annotations
|
| 8 |
+
|
| 9 |
+
from datetime import datetime, timezone
|
| 10 |
+
|
| 11 |
+
from fastapi import APIRouter, Depends, File, Form, HTTPException, Request, UploadFile, status
|
| 12 |
+
|
| 13 |
+
from app.core.config import BackendSettings
|
| 14 |
+
from app.core.logging import current_request_id
|
| 15 |
+
from app.schemas.caption import CaptionResponse, ErrorResponse, HealthResponse
|
| 16 |
+
from app.services.predictor_service import (
|
| 17 |
+
DECODE_STRATEGIES,
|
| 18 |
+
PredictorBusyError,
|
| 19 |
+
PredictorService,
|
| 20 |
+
)
|
| 21 |
+
from app.utils.image import ALLOWED_CONTENT_TYPES, ImageDecodeError, ImageTooLargeError
|
| 22 |
+
from captioning.utils import get_logger
|
| 23 |
+
|
| 24 |
+
log = get_logger(__name__)
|
| 25 |
+
|
| 26 |
+
router = APIRouter()
|
| 27 |
+
|
| 28 |
+
# The overload 503. One caption took 0.8-0.9 s of inference in the deploy smoke
|
| 29 |
+
# tests, so a slot usually frees within a couple of seconds.
|
| 30 |
+
BUSY_DETAIL = "The server is busy with other captions. Retry in a few seconds."
|
| 31 |
+
BUSY_RETRY_AFTER_SECONDS = 2
|
| 32 |
+
|
| 33 |
+
|
| 34 |
+
async def get_predictor_service(request: Request) -> PredictorService:
|
| 35 |
+
"""Resolve the singleton ``PredictorService`` from app state.
|
| 36 |
+
|
| 37 |
+
Returns 503 instead of crashing if the lifespan hasn't finished loading
|
| 38 |
+
weights yet (which can happen if ``/v1/captions`` is hit during a
|
| 39 |
+
rolling restart). Async, so FastAPI resolves it on the event loop rather
|
| 40 |
+
than in a worker thread that busy inference might be holding.
|
| 41 |
+
"""
|
| 42 |
+
service: PredictorService | None = getattr(request.app.state, "predictor_service", None)
|
| 43 |
+
if service is None:
|
| 44 |
+
raise HTTPException(
|
| 45 |
+
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
|
| 46 |
+
detail="Predictor is not ready yet.",
|
| 47 |
+
)
|
| 48 |
+
return service
|
| 49 |
+
|
| 50 |
+
|
| 51 |
+
@router.get(
|
| 52 |
+
"/healthz",
|
| 53 |
+
response_model=HealthResponse,
|
| 54 |
+
tags=["health"],
|
| 55 |
+
summary="Liveness + readiness probe",
|
| 56 |
+
)
|
| 57 |
+
async def healthz(request: Request) -> HealthResponse:
|
| 58 |
+
"""Return readiness state. Always 200 — readiness is conveyed by ``model_loaded``.
|
| 59 |
+
|
| 60 |
+
Runs entirely on the event loop, with no sync dependency, so it never
|
| 61 |
+
waits for a worker thread that inference might be holding.
|
| 62 |
+
"""
|
| 63 |
+
settings: BackendSettings = request.app.state.backend_settings
|
| 64 |
+
service: PredictorService | None = getattr(request.app.state, "predictor_service", None)
|
| 65 |
+
return HealthResponse(
|
| 66 |
+
status="ok" if service is not None else "loading",
|
| 67 |
+
model_loaded=service is not None,
|
| 68 |
+
model_version=service.model_version if service is not None else settings.model_version,
|
| 69 |
+
api_version=settings.api_version,
|
| 70 |
+
timestamp=datetime.now(timezone.utc),
|
| 71 |
+
)
|
| 72 |
+
|
| 73 |
+
|
| 74 |
+
@router.post(
|
| 75 |
+
"/v1/captions",
|
| 76 |
+
response_model=CaptionResponse,
|
| 77 |
+
tags=["captions"],
|
| 78 |
+
status_code=status.HTTP_200_OK,
|
| 79 |
+
summary="Generate a caption for an uploaded image",
|
| 80 |
+
responses={
|
| 81 |
+
400: {"model": ErrorResponse, "description": "Empty upload."},
|
| 82 |
+
413: {
|
| 83 |
+
"model": ErrorResponse,
|
| 84 |
+
"description": "Image exceeds the upload size or pixel-dimension limit.",
|
| 85 |
+
},
|
| 86 |
+
415: {"model": ErrorResponse, "description": "Unsupported image content type."},
|
| 87 |
+
422: {
|
| 88 |
+
"model": ErrorResponse,
|
| 89 |
+
"description": "Image bytes could not be decoded, or `decode_strategy` is unknown.",
|
| 90 |
+
},
|
| 91 |
+
503: {
|
| 92 |
+
"model": ErrorResponse,
|
| 93 |
+
"description": "Predictor not ready, or every inference slot is busy (with `Retry-After`).",
|
| 94 |
+
},
|
| 95 |
+
},
|
| 96 |
+
)
|
| 97 |
+
async def caption_image(
|
| 98 |
+
image: UploadFile = File(
|
| 99 |
+
...,
|
| 100 |
+
description="Image file to caption. Allowed: JPEG, PNG, BMP.",
|
| 101 |
+
),
|
| 102 |
+
decode_strategy: str | None = Form(
|
| 103 |
+
None,
|
| 104 |
+
description=(
|
| 105 |
+
"Decoding strategy for this request: `greedy` (the default) or `beam`. "
|
| 106 |
+
"Beam search uses the server's configured width and penalties."
|
| 107 |
+
),
|
| 108 |
+
),
|
| 109 |
+
service: PredictorService = Depends(get_predictor_service),
|
| 110 |
+
) -> CaptionResponse:
|
| 111 |
+
"""Accept a multipart image upload and return a generated caption."""
|
| 112 |
+
if decode_strategy is not None and decode_strategy not in DECODE_STRATEGIES:
|
| 113 |
+
raise HTTPException(
|
| 114 |
+
status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
|
| 115 |
+
detail=(
|
| 116 |
+
f"Unsupported decode_strategy: {decode_strategy!r}. "
|
| 117 |
+
f"Allowed: {list(DECODE_STRATEGIES)}."
|
| 118 |
+
),
|
| 119 |
+
)
|
| 120 |
+
if image.content_type not in ALLOWED_CONTENT_TYPES:
|
| 121 |
+
raise HTTPException(
|
| 122 |
+
status_code=status.HTTP_415_UNSUPPORTED_MEDIA_TYPE,
|
| 123 |
+
detail=(
|
| 124 |
+
f"Unsupported content type: {image.content_type!r}. "
|
| 125 |
+
f"Allowed: {sorted(ALLOWED_CONTENT_TYPES)}."
|
| 126 |
+
),
|
| 127 |
+
)
|
| 128 |
+
|
| 129 |
+
# Read at most one byte past the limit: enough to detect an oversized
|
| 130 |
+
# upload without pulling the whole file into memory.
|
| 131 |
+
payload = await image.read(service.max_upload_bytes + 1)
|
| 132 |
+
if not payload:
|
| 133 |
+
raise HTTPException(
|
| 134 |
+
status_code=status.HTTP_400_BAD_REQUEST,
|
| 135 |
+
detail="Empty file upload.",
|
| 136 |
+
)
|
| 137 |
+
if len(payload) > service.max_upload_bytes:
|
| 138 |
+
raise HTTPException(
|
| 139 |
+
status_code=status.HTTP_413_CONTENT_TOO_LARGE,
|
| 140 |
+
detail=f"Image exceeds the {service.max_upload_bytes}-byte upload limit.",
|
| 141 |
+
)
|
| 142 |
+
|
| 143 |
+
try:
|
| 144 |
+
caption, latency_ms = await service.caption_image_bytes(
|
| 145 |
+
payload, decode_strategy=decode_strategy
|
| 146 |
+
)
|
| 147 |
+
except PredictorBusyError as exc:
|
| 148 |
+
raise HTTPException(
|
| 149 |
+
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
|
| 150 |
+
detail=BUSY_DETAIL,
|
| 151 |
+
headers={"Retry-After": str(BUSY_RETRY_AFTER_SECONDS)},
|
| 152 |
+
) from exc
|
| 153 |
+
except ImageTooLargeError as exc:
|
| 154 |
+
raise HTTPException(
|
| 155 |
+
status_code=status.HTTP_413_CONTENT_TOO_LARGE,
|
| 156 |
+
detail=str(exc),
|
| 157 |
+
) from exc
|
| 158 |
+
except ImageDecodeError as exc:
|
| 159 |
+
raise HTTPException(
|
| 160 |
+
status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
|
| 161 |
+
detail=str(exc),
|
| 162 |
+
) from exc
|
| 163 |
+
|
| 164 |
+
return CaptionResponse(
|
| 165 |
+
caption=caption,
|
| 166 |
+
model_version=service.model_version,
|
| 167 |
+
decode_strategy=decode_strategy or service.decode_strategy,
|
| 168 |
+
latency_ms=round(latency_ms, 2),
|
| 169 |
+
request_id=current_request_id(),
|
| 170 |
+
)
|
backend/app/core/__init__.py
ADDED
|
@@ -0,0 +1,18 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Core: backend settings and HTTP-layer logging glue."""
|
| 2 |
+
|
| 3 |
+
from app.core.config import BackendSettings, get_backend_settings
|
| 4 |
+
from app.core.logging import (
|
| 5 |
+
REQUEST_ID_HEADER,
|
| 6 |
+
RequestContextMiddleware,
|
| 7 |
+
configure_app_logging,
|
| 8 |
+
current_request_id,
|
| 9 |
+
)
|
| 10 |
+
|
| 11 |
+
__all__ = [
|
| 12 |
+
"REQUEST_ID_HEADER",
|
| 13 |
+
"BackendSettings",
|
| 14 |
+
"RequestContextMiddleware",
|
| 15 |
+
"configure_app_logging",
|
| 16 |
+
"current_request_id",
|
| 17 |
+
"get_backend_settings",
|
| 18 |
+
]
|
backend/app/core/body_limit.py
ADDED
|
@@ -0,0 +1,71 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Request-body size cap, enforced while the body is still arriving.
|
| 2 |
+
|
| 3 |
+
FastAPI parses a multipart form before any route or dependency runs, and
|
| 4 |
+
Starlette's parser reads the whole body, spooling file parts to a temporary
|
| 5 |
+
file past 1 MiB. The route's own size check therefore runs only after an
|
| 6 |
+
oversized upload has been received in full. This middleware rejects such a
|
| 7 |
+
body first: as soon as the declared ``Content-Length``, or the bytes received
|
| 8 |
+
so far, exceed the cap.
|
| 9 |
+
|
| 10 |
+
It raises the same 413 ``HTTPException`` the route uses, from inside
|
| 11 |
+
``receive()``. FastAPI re-raises an ``HTTPException`` that occurs while it
|
| 12 |
+
reads the body, so the response keeps the route's ``{"detail": ...}`` shape
|
| 13 |
+
and passes through the CORS and request-id middleware like any other error.
|
| 14 |
+
"""
|
| 15 |
+
|
| 16 |
+
from __future__ import annotations
|
| 17 |
+
|
| 18 |
+
from starlette.exceptions import HTTPException
|
| 19 |
+
from starlette.status import HTTP_413_CONTENT_TOO_LARGE
|
| 20 |
+
from starlette.types import ASGIApp, Message, Receive, Scope, Send
|
| 21 |
+
|
| 22 |
+
# Room for the multipart framing around the image: boundary lines and the
|
| 23 |
+
# part's headers. The route still enforces the exact image-size limit.
|
| 24 |
+
MULTIPART_OVERHEAD_BYTES = 64 * 1024
|
| 25 |
+
|
| 26 |
+
|
| 27 |
+
class BodySizeLimitMiddleware:
|
| 28 |
+
"""Reject request bodies larger than ``max_upload_bytes`` plus framing."""
|
| 29 |
+
|
| 30 |
+
def __init__(self, app: ASGIApp, *, max_upload_bytes: int) -> None:
|
| 31 |
+
self.app = app
|
| 32 |
+
self.max_upload_bytes = max_upload_bytes
|
| 33 |
+
self.max_body_bytes = max_upload_bytes + MULTIPART_OVERHEAD_BYTES
|
| 34 |
+
|
| 35 |
+
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
|
| 36 |
+
if scope["type"] != "http":
|
| 37 |
+
await self.app(scope, receive, send)
|
| 38 |
+
return
|
| 39 |
+
|
| 40 |
+
declared = _declared_length(scope)
|
| 41 |
+
received = 0
|
| 42 |
+
|
| 43 |
+
async def limited_receive() -> Message:
|
| 44 |
+
nonlocal received
|
| 45 |
+
if declared is not None and declared > self.max_body_bytes:
|
| 46 |
+
raise self._too_large()
|
| 47 |
+
message = await receive()
|
| 48 |
+
if message["type"] == "http.request":
|
| 49 |
+
received += len(message.get("body", b""))
|
| 50 |
+
if received > self.max_body_bytes:
|
| 51 |
+
raise self._too_large()
|
| 52 |
+
return message
|
| 53 |
+
|
| 54 |
+
await self.app(scope, limited_receive, send)
|
| 55 |
+
|
| 56 |
+
def _too_large(self) -> HTTPException:
|
| 57 |
+
return HTTPException(
|
| 58 |
+
status_code=HTTP_413_CONTENT_TOO_LARGE,
|
| 59 |
+
detail=f"Image exceeds the {self.max_upload_bytes}-byte upload limit.",
|
| 60 |
+
)
|
| 61 |
+
|
| 62 |
+
|
| 63 |
+
def _declared_length(scope: Scope) -> int | None:
|
| 64 |
+
"""Return the request's ``Content-Length``, or ``None`` if absent or invalid."""
|
| 65 |
+
for name, value in scope["headers"]:
|
| 66 |
+
if name == b"content-length":
|
| 67 |
+
try:
|
| 68 |
+
return int(value)
|
| 69 |
+
except ValueError:
|
| 70 |
+
return None
|
| 71 |
+
return None
|
backend/app/core/config.py
ADDED
|
@@ -0,0 +1,120 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Backend runtime settings.
|
| 2 |
+
|
| 3 |
+
These settings drive the FastAPI process itself: where to find the trained
|
| 4 |
+
artifacts, what to advertise as the model version, whether to warm up at
|
| 5 |
+
boot. They are intentionally separate from ``captioning.config.AppConfig``,
|
| 6 |
+
which owns the *ML* configuration (architecture, decode strategy, CORS
|
| 7 |
+
origins). Keeping the two layers split lets ops change deployment paths
|
| 8 |
+
without touching research configs, and vice versa.
|
| 9 |
+
|
| 10 |
+
Override any field via environment variable, prefixed with ``BACKEND_``::
|
| 11 |
+
|
| 12 |
+
BACKEND_CONFIG_PATH=configs/base.yaml
|
| 13 |
+
BACKEND_WEIGHTS_PATH=models/v1.0.0/model.h5
|
| 14 |
+
BACKEND_TOKENIZER_DIR=models/v1.0.0
|
| 15 |
+
BACKEND_MODEL_VERSION=v1.0.0
|
| 16 |
+
BACKEND_WARMUP=true
|
| 17 |
+
BACKEND_WEIGHTS_HUB_REPO=your-username/captioning-weights
|
| 18 |
+
BACKEND_WEIGHTS_HUB_REVISION=<full 40-character commit id>
|
| 19 |
+
"""
|
| 20 |
+
|
| 21 |
+
from __future__ import annotations
|
| 22 |
+
|
| 23 |
+
from functools import lru_cache
|
| 24 |
+
from pathlib import Path
|
| 25 |
+
|
| 26 |
+
from pydantic import Field, field_validator
|
| 27 |
+
from pydantic_settings import BaseSettings, SettingsConfigDict
|
| 28 |
+
|
| 29 |
+
|
| 30 |
+
class BackendSettings(BaseSettings):
|
| 31 |
+
"""Settings for the FastAPI inference service."""
|
| 32 |
+
|
| 33 |
+
config_path: Path = Field(
|
| 34 |
+
default=Path("configs/base.yaml"),
|
| 35 |
+
description="Path to the YAML AppConfig consumed by the ML package.",
|
| 36 |
+
)
|
| 37 |
+
weights_path: Path = Field(
|
| 38 |
+
default=Path("models/v1.0.0/model.h5"),
|
| 39 |
+
description="Path to the trained Keras weights file (used when weights_hub_repo is unset).",
|
| 40 |
+
)
|
| 41 |
+
tokenizer_dir: Path = Field(
|
| 42 |
+
default=Path("models/v1.0.0"),
|
| 43 |
+
description=(
|
| 44 |
+
"Directory containing vocab.json (or a legacy vocab.pkl) "
|
| 45 |
+
"(used when weights_hub_repo is unset)."
|
| 46 |
+
),
|
| 47 |
+
)
|
| 48 |
+
model_version: str = Field(
|
| 49 |
+
default="v1.0.0",
|
| 50 |
+
description="Semantic version surfaced in /healthz and caption responses.",
|
| 51 |
+
)
|
| 52 |
+
api_version: str = Field(
|
| 53 |
+
default="0.1.0",
|
| 54 |
+
description="FastAPI app version (shown in OpenAPI docs).",
|
| 55 |
+
)
|
| 56 |
+
warmup: bool = Field(
|
| 57 |
+
default=True,
|
| 58 |
+
description="Run one dummy inference at startup so the first request is fast.",
|
| 59 |
+
)
|
| 60 |
+
request_id_header: str = Field(
|
| 61 |
+
default="x-request-id",
|
| 62 |
+
description="HTTP header used for request correlation IDs.",
|
| 63 |
+
)
|
| 64 |
+
|
| 65 |
+
# ---- HuggingFace Hub weights pull (WS-A4) -------------------------------
|
| 66 |
+
# When ``weights_hub_repo`` is set, ``lifespan`` calls
|
| 67 |
+
# ``huggingface_hub.snapshot_download`` for the weights file and vocab.json
|
| 68 |
+
# at an immutable commit, and resolves ``weights_path`` and
|
| 69 |
+
# ``tokenizer_dir`` to paths inside the downloaded snapshot. This lets the
|
| 70 |
+
# Docker image stay small and lets us rotate weights without rebuilding.
|
| 71 |
+
weights_hub_repo: str | None = Field(
|
| 72 |
+
default=None,
|
| 73 |
+
description="HuggingFace Hub repo id (e.g. 'user/captioning-weights'). None = use local paths.",
|
| 74 |
+
)
|
| 75 |
+
weights_hub_revision: str = Field(
|
| 76 |
+
default="main",
|
| 77 |
+
description=(
|
| 78 |
+
"Full 40-character Hub commit id, or a tag pinned in "
|
| 79 |
+
"app.services.weights_loader.PINNED_HUB_REVISIONS. Any other value, "
|
| 80 |
+
"including this 'main' default, is refused when weights_hub_repo is set, "
|
| 81 |
+
"and so is a commit without SHA-256 digests in PINNED_ARTIFACT_SHA256."
|
| 82 |
+
),
|
| 83 |
+
)
|
| 84 |
+
weights_hub_filename: str = Field(
|
| 85 |
+
default="model.h5",
|
| 86 |
+
description="Filename of the weights file inside the Hub snapshot.",
|
| 87 |
+
)
|
| 88 |
+
weights_cache_dir: Path | None = Field(
|
| 89 |
+
default=None,
|
| 90 |
+
description="Local cache dir for snapshot_download. None = HF Hub default ($HF_HOME).",
|
| 91 |
+
)
|
| 92 |
+
|
| 93 |
+
model_config = SettingsConfigDict(
|
| 94 |
+
env_prefix="BACKEND_",
|
| 95 |
+
case_sensitive=False,
|
| 96 |
+
extra="ignore",
|
| 97 |
+
# ``model_version`` would otherwise warn about pydantic's ``model_``
|
| 98 |
+
# namespace on every import; ``settings_`` stays protected.
|
| 99 |
+
protected_namespaces=("settings_",),
|
| 100 |
+
)
|
| 101 |
+
|
| 102 |
+
@field_validator("config_path", "weights_path", "tokenizer_dir")
|
| 103 |
+
@classmethod
|
| 104 |
+
def _expand_user(cls, value: Path) -> Path:
|
| 105 |
+
return value.expanduser()
|
| 106 |
+
|
| 107 |
+
@field_validator("weights_cache_dir")
|
| 108 |
+
@classmethod
|
| 109 |
+
def _expand_optional_user(cls, value: Path | None) -> Path | None:
|
| 110 |
+
return value.expanduser() if value is not None else None
|
| 111 |
+
|
| 112 |
+
|
| 113 |
+
@lru_cache(maxsize=1)
|
| 114 |
+
def get_backend_settings() -> BackendSettings:
|
| 115 |
+
"""Return a process-wide ``BackendSettings`` instance.
|
| 116 |
+
|
| 117 |
+
Cached so env-var parsing happens once. Tests that need to override env
|
| 118 |
+
can call ``get_backend_settings.cache_clear()`` between cases.
|
| 119 |
+
"""
|
| 120 |
+
return BackendSettings()
|
backend/app/core/logging.py
ADDED
|
@@ -0,0 +1,87 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""HTTP-layer logging glue.
|
| 2 |
+
|
| 3 |
+
The ML package already configures structlog (``captioning.utils.logging``).
|
| 4 |
+
The FastAPI process has two extra needs on top of that:
|
| 5 |
+
|
| 6 |
+
1. **Request correlation** — every log line emitted while handling a
|
| 7 |
+
request should carry the same ``request_id`` so logs can be grouped.
|
| 8 |
+
We bind it once via ``structlog.contextvars`` so any ``log.info(...)``
|
| 9 |
+
downstream automatically inherits it without threading the id through
|
| 10 |
+
function signatures.
|
| 11 |
+
|
| 12 |
+
2. **Access logs as structured events** — uvicorn's default access log is
|
| 13 |
+
a plain string. Re-emitting one structured ``request_finished`` event
|
| 14 |
+
per request keeps the log stream homogeneous and indexable.
|
| 15 |
+
"""
|
| 16 |
+
|
| 17 |
+
from __future__ import annotations
|
| 18 |
+
|
| 19 |
+
import time
|
| 20 |
+
import uuid
|
| 21 |
+
from collections.abc import Awaitable, Callable
|
| 22 |
+
|
| 23 |
+
import structlog
|
| 24 |
+
from starlette.middleware.base import BaseHTTPMiddleware
|
| 25 |
+
from starlette.requests import Request
|
| 26 |
+
from starlette.responses import Response
|
| 27 |
+
|
| 28 |
+
from captioning.utils.logging import configure_logging, get_logger
|
| 29 |
+
|
| 30 |
+
log = get_logger(__name__)
|
| 31 |
+
|
| 32 |
+
REQUEST_ID_HEADER = "x-request-id"
|
| 33 |
+
|
| 34 |
+
|
| 35 |
+
def configure_app_logging() -> None:
|
| 36 |
+
"""Initialise structlog for the FastAPI process.
|
| 37 |
+
|
| 38 |
+
Idempotent — delegates to the ML package's ``configure_logging`` so dev
|
| 39 |
+
gets pretty colourised output and ``APP_ENV=production`` flips to JSON.
|
| 40 |
+
"""
|
| 41 |
+
configure_logging()
|
| 42 |
+
|
| 43 |
+
|
| 44 |
+
class RequestContextMiddleware(BaseHTTPMiddleware):
|
| 45 |
+
"""Bind a request id to structlog and log start/finish events.
|
| 46 |
+
|
| 47 |
+
The id comes from the inbound ``x-request-id`` header when present
|
| 48 |
+
(so an upstream gateway can stitch traces), or a fresh ``uuid4`` hex
|
| 49 |
+
otherwise. Either way it's echoed back on the response.
|
| 50 |
+
"""
|
| 51 |
+
|
| 52 |
+
async def dispatch(
|
| 53 |
+
self,
|
| 54 |
+
request: Request,
|
| 55 |
+
call_next: Callable[[Request], Awaitable[Response]],
|
| 56 |
+
) -> Response:
|
| 57 |
+
request_id = request.headers.get(REQUEST_ID_HEADER) or uuid.uuid4().hex
|
| 58 |
+
|
| 59 |
+
structlog.contextvars.clear_contextvars()
|
| 60 |
+
structlog.contextvars.bind_contextvars(
|
| 61 |
+
request_id=request_id,
|
| 62 |
+
method=request.method,
|
| 63 |
+
path=request.url.path,
|
| 64 |
+
)
|
| 65 |
+
|
| 66 |
+
start = time.perf_counter()
|
| 67 |
+
log.info("request_started")
|
| 68 |
+
try:
|
| 69 |
+
response = await call_next(request)
|
| 70 |
+
except Exception:
|
| 71 |
+
duration_ms = (time.perf_counter() - start) * 1000
|
| 72 |
+
log.exception("request_failed", duration_ms=round(duration_ms, 2))
|
| 73 |
+
raise
|
| 74 |
+
|
| 75 |
+
duration_ms = (time.perf_counter() - start) * 1000
|
| 76 |
+
log.info(
|
| 77 |
+
"request_finished",
|
| 78 |
+
status=response.status_code,
|
| 79 |
+
duration_ms=round(duration_ms, 2),
|
| 80 |
+
)
|
| 81 |
+
response.headers[REQUEST_ID_HEADER] = request_id
|
| 82 |
+
return response
|
| 83 |
+
|
| 84 |
+
|
| 85 |
+
def current_request_id() -> str:
|
| 86 |
+
"""Return the request id bound to the current contextvars, or ``""``."""
|
| 87 |
+
return str(structlog.contextvars.get_contextvars().get("request_id", ""))
|
backend/app/main.py
ADDED
|
@@ -0,0 +1,110 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""FastAPI application entrypoint.
|
| 2 |
+
|
| 3 |
+
Run locally with::
|
| 4 |
+
|
| 5 |
+
uvicorn --app-dir backend app.main:app --host 0.0.0.0 --port 8000 --reload
|
| 6 |
+
|
| 7 |
+
Lifespan order:
|
| 8 |
+
1. Load YAML ``AppConfig`` (research-side hyperparameters).
|
| 9 |
+
2. Load weights + tokenizer into a ``CaptionPredictor`` singleton.
|
| 10 |
+
3. Optionally warmup so the first request doesn't pay TF's lazy build cost.
|
| 11 |
+
4. Wrap the predictor in a ``PredictorService`` and stash on app state.
|
| 12 |
+
|
| 13 |
+
The singleton lives on ``app.state.predictor_service``; routes pull it
|
| 14 |
+
through a ``Depends`` so tests can override the dependency cleanly.
|
| 15 |
+
"""
|
| 16 |
+
|
| 17 |
+
from __future__ import annotations
|
| 18 |
+
|
| 19 |
+
from collections.abc import AsyncIterator
|
| 20 |
+
from contextlib import asynccontextmanager
|
| 21 |
+
|
| 22 |
+
from fastapi import FastAPI
|
| 23 |
+
from fastapi.middleware.cors import CORSMiddleware
|
| 24 |
+
|
| 25 |
+
from app.api.routes import router
|
| 26 |
+
from app.core.body_limit import BodySizeLimitMiddleware
|
| 27 |
+
from app.core.config import BackendSettings, get_backend_settings
|
| 28 |
+
from app.core.logging import RequestContextMiddleware, configure_app_logging
|
| 29 |
+
from app.services.predictor_service import PredictorService
|
| 30 |
+
from app.services.weights_loader import resolve_weights
|
| 31 |
+
from captioning.config import load_config
|
| 32 |
+
from captioning.config.schema import AppConfig
|
| 33 |
+
from captioning.inference import CaptionPredictor
|
| 34 |
+
from captioning.utils import get_logger
|
| 35 |
+
|
| 36 |
+
log = get_logger(__name__)
|
| 37 |
+
|
| 38 |
+
|
| 39 |
+
@asynccontextmanager
|
| 40 |
+
async def lifespan(app: FastAPI) -> AsyncIterator[None]:
|
| 41 |
+
"""Load the predictor at startup, release it at shutdown."""
|
| 42 |
+
settings: BackendSettings = app.state.backend_settings
|
| 43 |
+
config: AppConfig = app.state.app_config
|
| 44 |
+
|
| 45 |
+
weights_path, tokenizer_dir = resolve_weights(settings)
|
| 46 |
+
|
| 47 |
+
log.info(
|
| 48 |
+
"predictor_loading",
|
| 49 |
+
weights=str(weights_path),
|
| 50 |
+
tokenizer_dir=str(tokenizer_dir),
|
| 51 |
+
model_version=settings.model_version,
|
| 52 |
+
)
|
| 53 |
+
|
| 54 |
+
predictor = CaptionPredictor.from_artifacts(
|
| 55 |
+
weights_path=weights_path,
|
| 56 |
+
tokenizer_dir=tokenizer_dir,
|
| 57 |
+
config=config,
|
| 58 |
+
)
|
| 59 |
+
if settings.warmup:
|
| 60 |
+
predictor.warmup()
|
| 61 |
+
|
| 62 |
+
app.state.predictor_service = PredictorService(
|
| 63 |
+
predictor=predictor,
|
| 64 |
+
model_version=settings.model_version,
|
| 65 |
+
max_upload_bytes=config.serve.max_upload_bytes,
|
| 66 |
+
)
|
| 67 |
+
log.info("predictor_ready", model_version=settings.model_version)
|
| 68 |
+
|
| 69 |
+
try:
|
| 70 |
+
yield
|
| 71 |
+
finally:
|
| 72 |
+
app.state.predictor_service = None
|
| 73 |
+
log.info("predictor_unloaded")
|
| 74 |
+
|
| 75 |
+
|
| 76 |
+
def create_app() -> FastAPI:
|
| 77 |
+
"""Build the FastAPI app. Factory form so tests can construct fresh apps."""
|
| 78 |
+
configure_app_logging()
|
| 79 |
+
settings = get_backend_settings()
|
| 80 |
+
config = load_config(settings.config_path)
|
| 81 |
+
|
| 82 |
+
app = FastAPI(
|
| 83 |
+
title="Image Captioning API",
|
| 84 |
+
version=settings.api_version,
|
| 85 |
+
description=(
|
| 86 |
+
"Production-grade inference service for the IEEE-published "
|
| 87 |
+
"CNN+Transformer image captioning model."
|
| 88 |
+
),
|
| 89 |
+
lifespan=lifespan,
|
| 90 |
+
)
|
| 91 |
+
|
| 92 |
+
app.state.backend_settings = settings
|
| 93 |
+
app.state.app_config = config
|
| 94 |
+
app.state.predictor_service = None
|
| 95 |
+
|
| 96 |
+
app.add_middleware(BodySizeLimitMiddleware, max_upload_bytes=config.serve.max_upload_bytes)
|
| 97 |
+
app.add_middleware(
|
| 98 |
+
CORSMiddleware,
|
| 99 |
+
allow_origins=config.serve.cors_allowed_origins,
|
| 100 |
+
allow_methods=["GET", "POST", "OPTIONS"],
|
| 101 |
+
allow_headers=["*"],
|
| 102 |
+
allow_credentials=False,
|
| 103 |
+
)
|
| 104 |
+
app.add_middleware(RequestContextMiddleware)
|
| 105 |
+
|
| 106 |
+
app.include_router(router)
|
| 107 |
+
return app
|
| 108 |
+
|
| 109 |
+
|
| 110 |
+
app = create_app()
|
backend/app/schemas/__init__.py
ADDED
|
@@ -0,0 +1,5 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Pydantic request/response schemas."""
|
| 2 |
+
|
| 3 |
+
from app.schemas.caption import CaptionResponse, ErrorResponse, HealthResponse
|
| 4 |
+
|
| 5 |
+
__all__ = ["CaptionResponse", "ErrorResponse", "HealthResponse"]
|
backend/app/schemas/caption.py
ADDED
|
@@ -0,0 +1,65 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Pydantic request/response models for the captioning API.
|
| 2 |
+
|
| 3 |
+
Schemas live separately from routes so the OpenAPI spec is stable even
|
| 4 |
+
when handler logic changes. Every field is annotated with an example so
|
| 5 |
+
``/docs`` is self-explanatory to anyone reviewing the portfolio.
|
| 6 |
+
"""
|
| 7 |
+
|
| 8 |
+
from __future__ import annotations
|
| 9 |
+
|
| 10 |
+
from datetime import datetime
|
| 11 |
+
|
| 12 |
+
from pydantic import BaseModel, ConfigDict, Field
|
| 13 |
+
|
| 14 |
+
|
| 15 |
+
class HealthResponse(BaseModel):
|
| 16 |
+
"""Liveness + readiness payload for ``GET /healthz``."""
|
| 17 |
+
|
| 18 |
+
status: str = Field(..., description="``ok`` once the predictor is loaded.")
|
| 19 |
+
model_loaded: bool = Field(..., description="True after weights + tokenizer are in memory.")
|
| 20 |
+
model_version: str = Field(..., description="Semantic version of the served model.")
|
| 21 |
+
api_version: str = Field(..., description="Backend release version.")
|
| 22 |
+
timestamp: datetime = Field(..., description="Server time the response was built (UTC).")
|
| 23 |
+
|
| 24 |
+
model_config = ConfigDict(
|
| 25 |
+
protected_namespaces=(),
|
| 26 |
+
json_schema_extra={
|
| 27 |
+
"example": {
|
| 28 |
+
"status": "ok",
|
| 29 |
+
"model_loaded": True,
|
| 30 |
+
"model_version": "v1.0.0",
|
| 31 |
+
"api_version": "0.1.0",
|
| 32 |
+
"timestamp": "2026-05-09T12:00:00Z",
|
| 33 |
+
}
|
| 34 |
+
},
|
| 35 |
+
)
|
| 36 |
+
|
| 37 |
+
|
| 38 |
+
class CaptionResponse(BaseModel):
|
| 39 |
+
"""Successful response from ``POST /v1/captions``."""
|
| 40 |
+
|
| 41 |
+
caption: str = Field(..., description="Generated caption text (without start/end tokens).")
|
| 42 |
+
model_version: str = Field(..., description="Model version that produced this caption.")
|
| 43 |
+
decode_strategy: str = Field(..., description="Decoding strategy used (e.g. ``greedy``).")
|
| 44 |
+
latency_ms: float = Field(..., description="Inference time in milliseconds.")
|
| 45 |
+
request_id: str = Field(..., description="Correlation id; matches the ``x-request-id`` header.")
|
| 46 |
+
|
| 47 |
+
model_config = ConfigDict(
|
| 48 |
+
protected_namespaces=(),
|
| 49 |
+
json_schema_extra={
|
| 50 |
+
"example": {
|
| 51 |
+
"caption": "a man riding a surfboard on a wave",
|
| 52 |
+
"model_version": "v1.0.0",
|
| 53 |
+
"decode_strategy": "greedy",
|
| 54 |
+
"latency_ms": 187.42,
|
| 55 |
+
"request_id": "8f1c2e3b4d5a4f8e9b0c1d2e3f4a5b6c",
|
| 56 |
+
}
|
| 57 |
+
},
|
| 58 |
+
)
|
| 59 |
+
|
| 60 |
+
|
| 61 |
+
class ErrorResponse(BaseModel):
|
| 62 |
+
"""Uniform error envelope returned by every non-2xx status."""
|
| 63 |
+
|
| 64 |
+
detail: str = Field(..., description="Human-readable error message.")
|
| 65 |
+
request_id: str = Field(default="", description="Correlation id for log lookup.")
|
backend/app/services/__init__.py
ADDED
|
@@ -0,0 +1,5 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Service layer wrapping the ML predictor."""
|
| 2 |
+
|
| 3 |
+
from app.services.predictor_service import PredictorService
|
| 4 |
+
|
| 5 |
+
__all__ = ["PredictorService"]
|
backend/app/services/predictor_service.py
ADDED
|
@@ -0,0 +1,145 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Service layer wrapping the ML ``CaptionPredictor``.
|
| 2 |
+
|
| 3 |
+
Why this exists between the route and the predictor:
|
| 4 |
+
* **Off-loop execution** — TensorFlow inference is sync and CPU-bound.
|
| 5 |
+
Running it inline blocks the event loop, so requests queue up
|
| 6 |
+
sequentially and event-loop-bound work (CORS, metrics, /healthz)
|
| 7 |
+
stalls. We push the call to a worker thread via ``anyio.to_thread``.
|
| 8 |
+
* **Bounded admission** — at most ``max_concurrent_inferences`` requests
|
| 9 |
+
decode or predict at once. The next one is refused straight away
|
| 10 |
+
(``PredictorBusyError``, a 503 at the route) instead of waiting.
|
| 11 |
+
* **Stable seam for testing** — routes depend on this class, not on
|
| 12 |
+
the concrete predictor. Tests can substitute a stub service that
|
| 13 |
+
returns canned captions without loading TensorFlow.
|
| 14 |
+
* **Future extension point** — Phase 4 will add a request batcher and
|
| 15 |
+
per-model registry behind the same ``caption_image_bytes`` API.
|
| 16 |
+
|
| 17 |
+
This class never re-implements inference; it delegates entirely to the
|
| 18 |
+
existing ``CaptionPredictor`` abstraction.
|
| 19 |
+
"""
|
| 20 |
+
|
| 21 |
+
from __future__ import annotations
|
| 22 |
+
|
| 23 |
+
import time
|
| 24 |
+
from functools import partial
|
| 25 |
+
|
| 26 |
+
from anyio import to_thread
|
| 27 |
+
|
| 28 |
+
from app.utils.image import bytes_to_tensor
|
| 29 |
+
from captioning.inference import CaptionPredictor
|
| 30 |
+
from captioning.inference.predictor import DECODE_STRATEGIES
|
| 31 |
+
from captioning.utils import get_logger
|
| 32 |
+
|
| 33 |
+
log = get_logger(__name__)
|
| 34 |
+
|
| 35 |
+
__all__ = [
|
| 36 |
+
"DECODE_STRATEGIES",
|
| 37 |
+
"MAX_CONCURRENT_INFERENCES",
|
| 38 |
+
"PredictorBusyError",
|
| 39 |
+
"PredictorService",
|
| 40 |
+
]
|
| 41 |
+
|
| 42 |
+
# Requests allowed to decode or predict at once. The Space runs on HF's
|
| 43 |
+
# cpu-basic tier (2 vCPU), and TensorFlow already spreads one inference across
|
| 44 |
+
# both cores, so more parallel inferences would compete for CPU rather than add
|
| 45 |
+
# throughput. Two also bounds the decode working set to about 1.2 GB by the
|
| 46 |
+
# TASK-026 estimate (~600 MB per image at the pixel limit). Chosen from the
|
| 47 |
+
# hardware, not benchmarked.
|
| 48 |
+
MAX_CONCURRENT_INFERENCES = 2
|
| 49 |
+
|
| 50 |
+
|
| 51 |
+
class PredictorBusyError(RuntimeError):
|
| 52 |
+
"""Raised when every inference slot is in use. The request isn't queued."""
|
| 53 |
+
|
| 54 |
+
|
| 55 |
+
class PredictorService:
|
| 56 |
+
"""Holds the singleton predictor and exposes async inference."""
|
| 57 |
+
|
| 58 |
+
def __init__(
|
| 59 |
+
self,
|
| 60 |
+
*,
|
| 61 |
+
predictor: CaptionPredictor,
|
| 62 |
+
model_version: str,
|
| 63 |
+
max_upload_bytes: int,
|
| 64 |
+
max_concurrent_inferences: int = MAX_CONCURRENT_INFERENCES,
|
| 65 |
+
) -> None:
|
| 66 |
+
"""Args:
|
| 67 |
+
predictor: A ready ``CaptionPredictor`` (weights already loaded).
|
| 68 |
+
model_version: Semver string surfaced in responses & health.
|
| 69 |
+
max_upload_bytes: Hard cap enforced at the route layer.
|
| 70 |
+
max_concurrent_inferences: Inference slots; at least 1.
|
| 71 |
+
"""
|
| 72 |
+
if max_concurrent_inferences < 1:
|
| 73 |
+
raise ValueError(
|
| 74 |
+
f"max_concurrent_inferences must be >= 1, got {max_concurrent_inferences}"
|
| 75 |
+
)
|
| 76 |
+
self._predictor = predictor
|
| 77 |
+
self._model_version = model_version
|
| 78 |
+
self._max_upload_bytes = max_upload_bytes
|
| 79 |
+
self._max_concurrent = max_concurrent_inferences
|
| 80 |
+
self._in_flight = 0
|
| 81 |
+
|
| 82 |
+
@property
|
| 83 |
+
def model_version(self) -> str:
|
| 84 |
+
return self._model_version
|
| 85 |
+
|
| 86 |
+
@property
|
| 87 |
+
def decode_strategy(self) -> str:
|
| 88 |
+
return self._predictor.decode_strategy
|
| 89 |
+
|
| 90 |
+
@property
|
| 91 |
+
def max_upload_bytes(self) -> int:
|
| 92 |
+
return self._max_upload_bytes
|
| 93 |
+
|
| 94 |
+
@property
|
| 95 |
+
def inferences_in_flight(self) -> int:
|
| 96 |
+
"""Requests holding a slot, including cancelled ones whose thread still runs."""
|
| 97 |
+
return self._in_flight
|
| 98 |
+
|
| 99 |
+
async def caption_image_bytes(
|
| 100 |
+
self, image_bytes: bytes, *, decode_strategy: str | None = None
|
| 101 |
+
) -> tuple[str, float]:
|
| 102 |
+
"""Decode bytes, run inference, and return (caption, latency_ms).
|
| 103 |
+
|
| 104 |
+
``decode_strategy`` applies to this call only; ``None`` means the
|
| 105 |
+
predictor's default. The route validates the value first.
|
| 106 |
+
|
| 107 |
+
Both the decode and the predict are offloaded to a worker thread so
|
| 108 |
+
the event loop stays responsive. Latency is measured around the
|
| 109 |
+
predict call only — decode timing belongs to a separate span if we
|
| 110 |
+
ever need it.
|
| 111 |
+
|
| 112 |
+
Admission: the slot check and the increment run on the event loop
|
| 113 |
+
with no ``await`` between them, so two requests can't both take the
|
| 114 |
+
last slot. ``finally`` gives the slot back on every exit path.
|
| 115 |
+
|
| 116 |
+
Cancellation: ``to_thread.run_sync`` keeps its default
|
| 117 |
+
``abandon_on_cancel=False``, so a cancelled caller (a client
|
| 118 |
+
disconnect, for instance) still waits for the worker thread it
|
| 119 |
+
started, and its slot is freed only when that thread has finished.
|
| 120 |
+
Cancelling doesn't stop TensorFlow work that is already running.
|
| 121 |
+
|
| 122 |
+
Raises:
|
| 123 |
+
PredictorBusyError: If every slot is in use. Nothing is decoded.
|
| 124 |
+
"""
|
| 125 |
+
if self._in_flight >= self._max_concurrent:
|
| 126 |
+
raise PredictorBusyError("All inference slots are in use.")
|
| 127 |
+
self._in_flight += 1
|
| 128 |
+
try:
|
| 129 |
+
tensor = await to_thread.run_sync(bytes_to_tensor, image_bytes)
|
| 130 |
+
|
| 131 |
+
start = time.perf_counter()
|
| 132 |
+
caption: str = await to_thread.run_sync(
|
| 133 |
+
partial(self._predictor.predict_tensor, tensor, decode_strategy=decode_strategy)
|
| 134 |
+
)
|
| 135 |
+
latency_ms = (time.perf_counter() - start) * 1000
|
| 136 |
+
finally:
|
| 137 |
+
self._in_flight -= 1
|
| 138 |
+
|
| 139 |
+
log.info(
|
| 140 |
+
"inference_completed",
|
| 141 |
+
model_version=self._model_version,
|
| 142 |
+
decode_strategy=decode_strategy or self.decode_strategy,
|
| 143 |
+
latency_ms=round(latency_ms, 2),
|
| 144 |
+
)
|
| 145 |
+
return caption, latency_ms
|
backend/app/services/weights_loader.py
ADDED
|
@@ -0,0 +1,257 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Resolve weights and tokenizer paths, optionally pulling from HuggingFace Hub.
|
| 2 |
+
|
| 3 |
+
In production we don't want to bake a 158 MB ``.h5`` file into the Docker
|
| 4 |
+
image — it makes builds slow and couples weight rotation to image rebuilds.
|
| 5 |
+
Instead, the image carries only code, and the runtime pulls the weights file
|
| 6 |
+
and ``vocab.json`` from a public Hub repo, at an immutable commit, the first
|
| 7 |
+
time the container boots. On HuggingFace Spaces the cache persists across
|
| 8 |
+
restarts. Both files are checked against SHA-256 digests pinned for that
|
| 9 |
+
commit before their paths are returned. The loaders reopen them by path, so
|
| 10 |
+
the check covers the cache as it was at that moment.
|
| 11 |
+
|
| 12 |
+
When ``BackendSettings.weights_hub_repo`` is unset we fall back to the
|
| 13 |
+
local paths declared in settings, which is what unit tests and `make serve`
|
| 14 |
+
use today. Local files aren't hashed.
|
| 15 |
+
"""
|
| 16 |
+
|
| 17 |
+
from __future__ import annotations
|
| 18 |
+
|
| 19 |
+
import hashlib
|
| 20 |
+
import re
|
| 21 |
+
from pathlib import Path
|
| 22 |
+
from typing import Protocol
|
| 23 |
+
|
| 24 |
+
from app.core.config import BackendSettings
|
| 25 |
+
from captioning.preprocessing.tokenizer import VOCAB_JSON_FILENAME
|
| 26 |
+
from captioning.utils import get_logger
|
| 27 |
+
|
| 28 |
+
log = get_logger(__name__)
|
| 29 |
+
|
| 30 |
+
# A full Hub commit id. Branches and tags can be moved to another commit; a
|
| 31 |
+
# commit id can't.
|
| 32 |
+
_COMMIT_ID = re.compile(r"[0-9a-f]{40}")
|
| 33 |
+
|
| 34 |
+
# Published tags and the commit each one named when it was evaluated and
|
| 35 |
+
# deployed (checked against the Hub API on 2026-10-12). The loader downloads the
|
| 36 |
+
# commit recorded here, never the tag, so moving a tag on the Hub changes
|
| 37 |
+
# nothing. A new tag needs a new line here before a Space variable can name it.
|
| 38 |
+
PINNED_HUB_REVISIONS: dict[tuple[str, str], str] = {
|
| 39 |
+
(
|
| 40 |
+
"apoorvrajdev/captioning-inceptionv3-transformer",
|
| 41 |
+
"v1.0.0",
|
| 42 |
+
): "d917a47decc1f352e50ba4226e92ca5f00d47e9a",
|
| 43 |
+
(
|
| 44 |
+
"apoorvrajdev/captioning-inceptionv3-transformer",
|
| 45 |
+
"v2.0.0",
|
| 46 |
+
): "59d93b4babb16b0ac81eef598f3abc271a355cbf",
|
| 47 |
+
}
|
| 48 |
+
|
| 49 |
+
# A lower-case hex SHA-256 digest, the form hashlib and Git LFS pointers use.
|
| 50 |
+
_SHA256 = re.compile(r"[0-9a-f]{64}")
|
| 51 |
+
|
| 52 |
+
# SHA-256 of every file serving reads, per repo and commit. A commit without an
|
| 53 |
+
# entry here, or a file missing from its entry, is refused before any download.
|
| 54 |
+
# Provenance (Hub API, read-only, 2026-10-12):
|
| 55 |
+
# - model.h5 is a Git LFS file. Its digest is the ``oid sha256`` of the LFS
|
| 56 |
+
# pointer stored in the commit's tree, and that pointer's git blob id matches
|
| 57 |
+
# the tree entry. v2.0.0's digest also appears in docs/EVAL_METHODOLOGY.md
|
| 58 |
+
# (2026-10-06).
|
| 59 |
+
# - vocab.json is a plain git file, for which the Hub publishes only a SHA-1
|
| 60 |
+
# blob id. Its digest is the SHA-256 of bytes whose git blob id matched the
|
| 61 |
+
# commit's tree entry (v1.0.0 f8ce580…, v2.0.0 1a437e1…). v1.0.0's also equals
|
| 62 |
+
# the repo's own models/v1.0.0/vocab.json.
|
| 63 |
+
# Serving a new commit means adding its digests here, from the same sources,
|
| 64 |
+
# and deploying that first.
|
| 65 |
+
PINNED_ARTIFACT_SHA256: dict[tuple[str, str], dict[str, str]] = {
|
| 66 |
+
(
|
| 67 |
+
"apoorvrajdev/captioning-inceptionv3-transformer",
|
| 68 |
+
"d917a47decc1f352e50ba4226e92ca5f00d47e9a", # v1.0.0
|
| 69 |
+
): {
|
| 70 |
+
"model.h5": "bfe020d920aa2f3d019bf7b5b33904384057372e7c304a9e101a2a59fe110084",
|
| 71 |
+
"vocab.json": "45ec1704d73046303cbd5292590b2e204b194a2d8345dfb84de81370b4ab4eef",
|
| 72 |
+
},
|
| 73 |
+
(
|
| 74 |
+
"apoorvrajdev/captioning-inceptionv3-transformer",
|
| 75 |
+
"59d93b4babb16b0ac81eef598f3abc271a355cbf", # v2.0.0
|
| 76 |
+
): {
|
| 77 |
+
"model.h5": "74963a3f7cd01b16f44cd179f1e21e9eb8d60b52d460e517cdcf38c015689b07",
|
| 78 |
+
"vocab.json": "4315c5a7350d6052dbc120e02d218b9a95a23797aa24ec5670824eb6667d2886",
|
| 79 |
+
},
|
| 80 |
+
}
|
| 81 |
+
|
| 82 |
+
|
| 83 |
+
class ArtifactIntegrityError(RuntimeError):
|
| 84 |
+
"""A downloaded file's SHA-256 differs from the digest pinned for its commit."""
|
| 85 |
+
|
| 86 |
+
|
| 87 |
+
class SnapshotDownloader(Protocol):
|
| 88 |
+
"""Minimal callable shape we need from ``huggingface_hub.snapshot_download``."""
|
| 89 |
+
|
| 90 |
+
def __call__(
|
| 91 |
+
self,
|
| 92 |
+
*,
|
| 93 |
+
repo_id: str,
|
| 94 |
+
revision: str,
|
| 95 |
+
cache_dir: str | None,
|
| 96 |
+
allow_patterns: list[str],
|
| 97 |
+
) -> str: ...
|
| 98 |
+
|
| 99 |
+
|
| 100 |
+
def resolve_hub_commit(repo_id: str, revision: str) -> str:
|
| 101 |
+
"""Return the immutable commit id to download for ``revision``.
|
| 102 |
+
|
| 103 |
+
A full 40-character commit id is used as given. A tag listed in
|
| 104 |
+
``PINNED_HUB_REVISIONS`` for ``repo_id`` resolves to its recorded commit.
|
| 105 |
+
Anything else, including a branch such as ``main``, raises ``ValueError``.
|
| 106 |
+
"""
|
| 107 |
+
if _COMMIT_ID.fullmatch(revision):
|
| 108 |
+
return revision
|
| 109 |
+
pinned = PINNED_HUB_REVISIONS.get((repo_id, revision))
|
| 110 |
+
if pinned is None:
|
| 111 |
+
raise ValueError(
|
| 112 |
+
f"BACKEND_WEIGHTS_HUB_REVISION={revision!r} is neither a full 40-character "
|
| 113 |
+
f"commit id nor a tag of {repo_id} pinned in "
|
| 114 |
+
"app.services.weights_loader.PINNED_HUB_REVISIONS. Set it to the commit id "
|
| 115 |
+
"(HfApi().model_info(repo_id, revision=tag).sha) or pin the tag first."
|
| 116 |
+
)
|
| 117 |
+
return pinned
|
| 118 |
+
|
| 119 |
+
|
| 120 |
+
def pinned_sha256(repo_id: str, commit: str, filenames: list[str]) -> dict[str, str]:
|
| 121 |
+
"""Return the pinned SHA-256 of each of ``filenames`` at ``commit``.
|
| 122 |
+
|
| 123 |
+
Raises ``ValueError`` when the commit has no entry in
|
| 124 |
+
``PINNED_ARTIFACT_SHA256``, when a file has no digest there, or when a
|
| 125 |
+
digest isn't 64 lower-case hex characters. Callers run this before
|
| 126 |
+
downloading, so an unverifiable revision never reaches the network.
|
| 127 |
+
"""
|
| 128 |
+
pinned = PINNED_ARTIFACT_SHA256.get((repo_id, commit))
|
| 129 |
+
if pinned is None:
|
| 130 |
+
raise ValueError(
|
| 131 |
+
f"Hub commit {commit} of {repo_id} has no SHA-256 digests in "
|
| 132 |
+
"app.services.weights_loader.PINNED_ARTIFACT_SHA256, so its files can't be "
|
| 133 |
+
"verified. Add the digests for that commit (see the comment above the table) "
|
| 134 |
+
"and deploy them before serving it."
|
| 135 |
+
)
|
| 136 |
+
expected: dict[str, str] = {}
|
| 137 |
+
for name in filenames:
|
| 138 |
+
digest = pinned.get(name)
|
| 139 |
+
if digest is None:
|
| 140 |
+
raise ValueError(
|
| 141 |
+
f"PINNED_ARTIFACT_SHA256 has no SHA-256 for {name} at Hub commit {commit} "
|
| 142 |
+
f"of {repo_id}. Pin a digest for every file serving reads."
|
| 143 |
+
)
|
| 144 |
+
if not _SHA256.fullmatch(digest):
|
| 145 |
+
raise ValueError(
|
| 146 |
+
f"PINNED_ARTIFACT_SHA256 holds a malformed SHA-256 for {name} at Hub commit "
|
| 147 |
+
f"{commit} of {repo_id}: expected 64 lower-case hex characters."
|
| 148 |
+
)
|
| 149 |
+
expected[name] = digest
|
| 150 |
+
return expected
|
| 151 |
+
|
| 152 |
+
|
| 153 |
+
def sha256_file(path: Path, chunk_size: int = 1 << 20) -> str:
|
| 154 |
+
"""Return the hex SHA-256 of ``path``, read in chunks so memory stays bounded."""
|
| 155 |
+
digest = hashlib.sha256()
|
| 156 |
+
with path.open("rb") as f:
|
| 157 |
+
while chunk := f.read(chunk_size):
|
| 158 |
+
digest.update(chunk)
|
| 159 |
+
return digest.hexdigest()
|
| 160 |
+
|
| 161 |
+
|
| 162 |
+
def verify_sha256(directory: Path, expected: dict[str, str], *, repo_id: str, commit: str) -> None:
|
| 163 |
+
"""Check every file in ``expected`` under ``directory`` against its digest.
|
| 164 |
+
|
| 165 |
+
Every file is hashed, even after a mismatch, so the error lists each file
|
| 166 |
+
that differs.
|
| 167 |
+
|
| 168 |
+
Raises:
|
| 169 |
+
ArtifactIntegrityError: One or more files differ from their digest.
|
| 170 |
+
"""
|
| 171 |
+
mismatches = []
|
| 172 |
+
for name, want in expected.items():
|
| 173 |
+
got = sha256_file(directory / name)
|
| 174 |
+
if got != want:
|
| 175 |
+
mismatches.append(f"{name} has SHA-256 {got}, expected {want}")
|
| 176 |
+
if mismatches:
|
| 177 |
+
raise ArtifactIntegrityError(
|
| 178 |
+
f"Hub commit {commit} of {repo_id} failed verification in {directory}: "
|
| 179 |
+
f"{'; '.join(mismatches)}. Nothing was loaded. The download or the local "
|
| 180 |
+
"cache is corrupt or was altered: delete the file from the Hub cache, both "
|
| 181 |
+
"the snapshot entry and the blob it links to, and restart to download it "
|
| 182 |
+
"again. If it still fails, the Hub no longer serves the pinned bytes; don't "
|
| 183 |
+
"serve this revision."
|
| 184 |
+
)
|
| 185 |
+
|
| 186 |
+
|
| 187 |
+
def resolve_weights(
|
| 188 |
+
settings: BackendSettings,
|
| 189 |
+
downloader: SnapshotDownloader | None = None,
|
| 190 |
+
) -> tuple[Path, Path]:
|
| 191 |
+
"""Return ``(weights_path, tokenizer_dir)`` for the predictor to load.
|
| 192 |
+
|
| 193 |
+
Local mode (``weights_hub_repo`` is None): returns the paths verbatim.
|
| 194 |
+
|
| 195 |
+
Hub mode: resolves the configured revision to a commit id, looks up the
|
| 196 |
+
SHA-256 pinned for each file at that commit, downloads only the weights
|
| 197 |
+
file and ``vocab.json`` at that commit, checks both against their digests,
|
| 198 |
+
then returns paths inside the snapshot directory. The downloader is
|
| 199 |
+
injectable so tests can substitute a stub instead of hitting the network.
|
| 200 |
+
|
| 201 |
+
Raises:
|
| 202 |
+
ValueError: The revision is not a commit id or a pinned tag, or the
|
| 203 |
+
commit has no valid pinned digest for a file.
|
| 204 |
+
FileNotFoundError: The commit lacks the weights file or ``vocab.json``.
|
| 205 |
+
ArtifactIntegrityError: A downloaded file differs from its digest.
|
| 206 |
+
"""
|
| 207 |
+
if not settings.weights_hub_repo:
|
| 208 |
+
log.info(
|
| 209 |
+
"weights_source_local",
|
| 210 |
+
weights=str(settings.weights_path),
|
| 211 |
+
tokenizer_dir=str(settings.tokenizer_dir),
|
| 212 |
+
)
|
| 213 |
+
return settings.weights_path, settings.tokenizer_dir
|
| 214 |
+
|
| 215 |
+
commit = resolve_hub_commit(settings.weights_hub_repo, settings.weights_hub_revision)
|
| 216 |
+
required = [settings.weights_hub_filename, VOCAB_JSON_FILENAME]
|
| 217 |
+
expected = pinned_sha256(settings.weights_hub_repo, commit, required)
|
| 218 |
+
|
| 219 |
+
if downloader is None:
|
| 220 |
+
from huggingface_hub import snapshot_download as _snapshot_download
|
| 221 |
+
|
| 222 |
+
downloader = _snapshot_download
|
| 223 |
+
assert downloader is not None # for type-checker
|
| 224 |
+
|
| 225 |
+
cache_dir = str(settings.weights_cache_dir) if settings.weights_cache_dir else None
|
| 226 |
+
log.info(
|
| 227 |
+
"weights_source_hub",
|
| 228 |
+
repo=settings.weights_hub_repo,
|
| 229 |
+
revision=settings.weights_hub_revision,
|
| 230 |
+
commit=commit,
|
| 231 |
+
files=required,
|
| 232 |
+
cache_dir=cache_dir,
|
| 233 |
+
)
|
| 234 |
+
snapshot_dir = Path(
|
| 235 |
+
downloader(
|
| 236 |
+
repo_id=settings.weights_hub_repo,
|
| 237 |
+
revision=commit,
|
| 238 |
+
cache_dir=cache_dir,
|
| 239 |
+
allow_patterns=required,
|
| 240 |
+
)
|
| 241 |
+
)
|
| 242 |
+
missing = [name for name in required if not (snapshot_dir / name).is_file()]
|
| 243 |
+
if missing:
|
| 244 |
+
raise FileNotFoundError(
|
| 245 |
+
f"Hub commit {commit} of {settings.weights_hub_repo} has no "
|
| 246 |
+
f"{', '.join(missing)}. Serving needs the weights file and "
|
| 247 |
+
f"{VOCAB_JSON_FILENAME} at the same commit."
|
| 248 |
+
)
|
| 249 |
+
verify_sha256(snapshot_dir, expected, repo_id=settings.weights_hub_repo, commit=commit)
|
| 250 |
+
weights_path = snapshot_dir / settings.weights_hub_filename
|
| 251 |
+
log.info(
|
| 252 |
+
"weights_verified",
|
| 253 |
+
snapshot_dir=str(snapshot_dir),
|
| 254 |
+
weights=str(weights_path),
|
| 255 |
+
sha256=expected,
|
| 256 |
+
)
|
| 257 |
+
return weights_path, snapshot_dir
|
backend/app/tests/__init__.py
ADDED
|
File without changes
|
backend/app/tests/conftest.py
ADDED
|
@@ -0,0 +1,101 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Shared fixtures for the backend test suite.
|
| 2 |
+
|
| 3 |
+
These tests deliberately avoid loading TensorFlow or any real model.
|
| 4 |
+
The route layer depends on ``PredictorService`` only through duck-typed
|
| 5 |
+
attributes (``model_version``, ``decode_strategy``, ``max_upload_bytes``,
|
| 6 |
+
``caption_image_bytes``), so a small fake stands in cleanly and keeps the
|
| 7 |
+
whole suite under one second.
|
| 8 |
+
|
| 9 |
+
We also bypass the FastAPI lifespan entirely. The lifespan builds a real
|
| 10 |
+
``CaptionPredictor`` from disk, which requires weights, a tokenizer, and a
|
| 11 |
+
TF graph build. Tests build a fresh ``FastAPI`` instance, wire the same
|
| 12 |
+
router and middleware, and stash the fake service directly on
|
| 13 |
+
``app.state.predictor_service`` — the exact slot the lifespan would have
|
| 14 |
+
populated in production.
|
| 15 |
+
"""
|
| 16 |
+
|
| 17 |
+
from __future__ import annotations
|
| 18 |
+
|
| 19 |
+
from collections.abc import Callable, Iterator
|
| 20 |
+
|
| 21 |
+
import pytest
|
| 22 |
+
from fastapi import FastAPI
|
| 23 |
+
from fastapi.testclient import TestClient
|
| 24 |
+
|
| 25 |
+
from app.api.routes import router
|
| 26 |
+
from app.core.config import BackendSettings
|
| 27 |
+
from app.core.logging import RequestContextMiddleware, configure_app_logging
|
| 28 |
+
from app.utils.image import ImageDecodeError, check_image_dimensions
|
| 29 |
+
|
| 30 |
+
configure_app_logging()
|
| 31 |
+
|
| 32 |
+
|
| 33 |
+
class FakePredictorService:
|
| 34 |
+
"""Duck-typed stand-in for ``PredictorService``."""
|
| 35 |
+
|
| 36 |
+
def __init__(
|
| 37 |
+
self,
|
| 38 |
+
*,
|
| 39 |
+
caption: str = "a test caption",
|
| 40 |
+
latency_ms: float = 1.23,
|
| 41 |
+
decode_strategy: str = "greedy",
|
| 42 |
+
model_version: str = "test-v0",
|
| 43 |
+
max_upload_bytes: int = 1024,
|
| 44 |
+
raise_decode_error: bool = False,
|
| 45 |
+
check_dimensions: bool = False,
|
| 46 |
+
) -> None:
|
| 47 |
+
self.model_version = model_version
|
| 48 |
+
self.decode_strategy = decode_strategy
|
| 49 |
+
self.max_upload_bytes = max_upload_bytes
|
| 50 |
+
self._caption = caption
|
| 51 |
+
self._latency_ms = latency_ms
|
| 52 |
+
self._raise = raise_decode_error
|
| 53 |
+
self._check_dimensions = check_dimensions
|
| 54 |
+
self.calls: list[bytes] = []
|
| 55 |
+
self.strategies: list[str | None] = []
|
| 56 |
+
|
| 57 |
+
async def caption_image_bytes(
|
| 58 |
+
self, image_bytes: bytes, *, decode_strategy: str | None = None
|
| 59 |
+
) -> tuple[str, float]:
|
| 60 |
+
self.calls.append(image_bytes)
|
| 61 |
+
self.strategies.append(decode_strategy)
|
| 62 |
+
if self._raise:
|
| 63 |
+
raise ImageDecodeError("synthetic decode failure")
|
| 64 |
+
if self._check_dimensions:
|
| 65 |
+
# The real header guard that ``bytes_to_tensor`` runs first, without TF.
|
| 66 |
+
check_image_dimensions(image_bytes)
|
| 67 |
+
return self._caption, self._latency_ms
|
| 68 |
+
|
| 69 |
+
|
| 70 |
+
def _build_app(service: FakePredictorService | None) -> FastAPI:
|
| 71 |
+
app = FastAPI()
|
| 72 |
+
app.state.backend_settings = BackendSettings()
|
| 73 |
+
app.state.predictor_service = service
|
| 74 |
+
app.add_middleware(RequestContextMiddleware)
|
| 75 |
+
app.include_router(router)
|
| 76 |
+
return app
|
| 77 |
+
|
| 78 |
+
|
| 79 |
+
@pytest.fixture
|
| 80 |
+
def fake_service() -> FakePredictorService:
|
| 81 |
+
return FakePredictorService()
|
| 82 |
+
|
| 83 |
+
|
| 84 |
+
@pytest.fixture
|
| 85 |
+
def client(fake_service: FakePredictorService) -> Iterator[TestClient]:
|
| 86 |
+
with TestClient(_build_app(fake_service)) as test_client:
|
| 87 |
+
yield test_client
|
| 88 |
+
|
| 89 |
+
|
| 90 |
+
@pytest.fixture
|
| 91 |
+
def client_without_service() -> Iterator[TestClient]:
|
| 92 |
+
with TestClient(_build_app(None)) as test_client:
|
| 93 |
+
yield test_client
|
| 94 |
+
|
| 95 |
+
|
| 96 |
+
@pytest.fixture
|
| 97 |
+
def build_client() -> Callable[[FakePredictorService | None], TestClient]:
|
| 98 |
+
def _make(service: FakePredictorService | None) -> TestClient:
|
| 99 |
+
return TestClient(_build_app(service))
|
| 100 |
+
|
| 101 |
+
return _make
|
backend/app/tests/image_bytes.py
ADDED
|
@@ -0,0 +1,47 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Small image byte builders for the upload dimension-guard tests.
|
| 2 |
+
|
| 3 |
+
A header-only image declares any size in a few dozen bytes, so tests can check
|
| 4 |
+
huge dimensions without allocating them. Nothing here decodes pixels.
|
| 5 |
+
"""
|
| 6 |
+
|
| 7 |
+
from __future__ import annotations
|
| 8 |
+
|
| 9 |
+
import io
|
| 10 |
+
import struct
|
| 11 |
+
import zlib
|
| 12 |
+
|
| 13 |
+
from PIL import Image
|
| 14 |
+
|
| 15 |
+
|
| 16 |
+
def encode(fmt: str, size: tuple[int, int] = (32, 24), mode: str = "RGB") -> bytes:
|
| 17 |
+
"""A small deterministic gradient image, saved as ``fmt``."""
|
| 18 |
+
width, height = size
|
| 19 |
+
pixels = bytes(
|
| 20 |
+
(x * 7 + y * 13 + channel * 50) % 256
|
| 21 |
+
for y in range(height)
|
| 22 |
+
for x in range(width)
|
| 23 |
+
for channel in range(3)
|
| 24 |
+
)
|
| 25 |
+
image = Image.frombytes("RGB", size, pixels).convert(mode)
|
| 26 |
+
buffer = io.BytesIO()
|
| 27 |
+
image.save(buffer, format=fmt)
|
| 28 |
+
return buffer.getvalue()
|
| 29 |
+
|
| 30 |
+
|
| 31 |
+
def png_header(width: int, height: int) -> bytes:
|
| 32 |
+
"""A structurally valid PNG that declares ``width`` x ``height`` and holds no pixels."""
|
| 33 |
+
|
| 34 |
+
def chunk(kind: bytes, data: bytes) -> bytes:
|
| 35 |
+
crc = zlib.crc32(kind + data)
|
| 36 |
+
return struct.pack(">I", len(data)) + kind + data + struct.pack(">I", crc)
|
| 37 |
+
|
| 38 |
+
ihdr = struct.pack(">IIBBBBB", width, height, 8, 2, 0, 0, 0) # 8-bit RGB
|
| 39 |
+
return b"\x89PNG\r\n\x1a\n" + chunk(b"IHDR", ihdr) + chunk(b"IDAT", b"") + chunk(b"IEND", b"")
|
| 40 |
+
|
| 41 |
+
|
| 42 |
+
def jpeg_header(width: int, height: int) -> bytes:
|
| 43 |
+
"""A baseline JPEG header (SOI, SOF0, SOS) for ``width`` x ``height``, with no scan data."""
|
| 44 |
+
components = b"\x01\x22\x00\x02\x11\x01\x03\x11\x01" # Y, Cb, Cr
|
| 45 |
+
sof0 = b"\xff\xc0" + struct.pack(">HBHHB", 17, 8, height, width, 3) + components
|
| 46 |
+
sos = b"\xff\xda" + struct.pack(">HB", 12, 3) + b"\x01\x00\x02\x11\x03\x11\x00\x3f\x00"
|
| 47 |
+
return b"\xff\xd8" + sof0 + sos
|
backend/app/tests/test_backend_settings.py
ADDED
|
@@ -0,0 +1,34 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""``BackendSettings`` imports cleanly and keeps its ``BACKEND_*`` field names."""
|
| 2 |
+
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
|
| 5 |
+
import subprocess
|
| 6 |
+
import sys
|
| 7 |
+
from pathlib import Path
|
| 8 |
+
|
| 9 |
+
import pytest
|
| 10 |
+
|
| 11 |
+
from app.core.config import BackendSettings
|
| 12 |
+
|
| 13 |
+
BACKEND_DIR = Path(__file__).resolve().parents[2]
|
| 14 |
+
|
| 15 |
+
|
| 16 |
+
def test_importing_the_settings_emits_no_user_warning() -> None:
|
| 17 |
+
"""A fresh interpreter, because the warning fires once, when the class is defined."""
|
| 18 |
+
result = subprocess.run(
|
| 19 |
+
[sys.executable, "-W", "error::UserWarning", "-c", "import app.core.config"],
|
| 20 |
+
cwd=BACKEND_DIR,
|
| 21 |
+
capture_output=True,
|
| 22 |
+
text=True,
|
| 23 |
+
check=False,
|
| 24 |
+
)
|
| 25 |
+
|
| 26 |
+
assert result.returncode == 0, result.stderr
|
| 27 |
+
|
| 28 |
+
|
| 29 |
+
def test_model_version_still_reads_its_environment_variable(
|
| 30 |
+
monkeypatch: pytest.MonkeyPatch,
|
| 31 |
+
) -> None:
|
| 32 |
+
monkeypatch.setenv("BACKEND_MODEL_VERSION", "v9.9.9")
|
| 33 |
+
|
| 34 |
+
assert BackendSettings().model_version == "v9.9.9"
|
backend/app/tests/test_body_size_limit.py
ADDED
|
@@ -0,0 +1,207 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Tests for ``BodySizeLimitMiddleware`` on ``POST /v1/captions``.
|
| 2 |
+
|
| 3 |
+
An oversized body must get the route's 413 before the multipart parser reads
|
| 4 |
+
it, whether or not the client declares ``Content-Length``. The app is driven
|
| 5 |
+
through raw ASGI, with the body handed over in 64 KiB chunks as uvicorn does,
|
| 6 |
+
so each test can count how much of the body the app actually read.
|
| 7 |
+
"""
|
| 8 |
+
|
| 9 |
+
from __future__ import annotations
|
| 10 |
+
|
| 11 |
+
import json
|
| 12 |
+
import tempfile
|
| 13 |
+
import weakref
|
| 14 |
+
from collections.abc import Callable
|
| 15 |
+
from typing import Any
|
| 16 |
+
|
| 17 |
+
import anyio
|
| 18 |
+
import pytest
|
| 19 |
+
from fastapi import FastAPI
|
| 20 |
+
from fastapi.middleware.cors import CORSMiddleware
|
| 21 |
+
from starlette.types import Message
|
| 22 |
+
|
| 23 |
+
from app.api.routes import router
|
| 24 |
+
from app.core.body_limit import MULTIPART_OVERHEAD_BYTES, BodySizeLimitMiddleware
|
| 25 |
+
from app.core.config import BackendSettings
|
| 26 |
+
from app.core.logging import REQUEST_ID_HEADER, RequestContextMiddleware
|
| 27 |
+
from app.tests.conftest import FakePredictorService
|
| 28 |
+
|
| 29 |
+
CHUNK = 64 * 1024
|
| 30 |
+
BOUNDARY = b"test-boundary"
|
| 31 |
+
ORIGIN = "http://localhost:5173"
|
| 32 |
+
|
| 33 |
+
|
| 34 |
+
def _build_app(service: FakePredictorService) -> FastAPI:
|
| 35 |
+
app = FastAPI()
|
| 36 |
+
app.state.backend_settings = BackendSettings()
|
| 37 |
+
app.state.predictor_service = service
|
| 38 |
+
# Same order as create_app: the cap innermost, then CORS, then request context.
|
| 39 |
+
app.add_middleware(BodySizeLimitMiddleware, max_upload_bytes=service.max_upload_bytes)
|
| 40 |
+
app.add_middleware(
|
| 41 |
+
CORSMiddleware,
|
| 42 |
+
allow_origins=[ORIGIN],
|
| 43 |
+
allow_methods=["GET", "POST", "OPTIONS"],
|
| 44 |
+
allow_headers=["*"],
|
| 45 |
+
allow_credentials=False,
|
| 46 |
+
)
|
| 47 |
+
app.add_middleware(RequestContextMiddleware)
|
| 48 |
+
app.include_router(router)
|
| 49 |
+
return app
|
| 50 |
+
|
| 51 |
+
|
| 52 |
+
def _multipart(image: bytes) -> bytes:
|
| 53 |
+
return (
|
| 54 |
+
b"--" + BOUNDARY + b"\r\n"
|
| 55 |
+
b'Content-Disposition: form-data; name="image"; filename="a.jpg"\r\n'
|
| 56 |
+
b"Content-Type: image/jpeg\r\n\r\n" + image + b"\r\n--" + BOUNDARY + b"--\r\n"
|
| 57 |
+
)
|
| 58 |
+
|
| 59 |
+
|
| 60 |
+
def _post(
|
| 61 |
+
app: FastAPI,
|
| 62 |
+
body: bytes,
|
| 63 |
+
*,
|
| 64 |
+
declare_length: bool,
|
| 65 |
+
after: Callable[[], None] | None = None,
|
| 66 |
+
) -> dict[str, Any]:
|
| 67 |
+
"""POST ``body`` to ``/v1/captions``; return the status, JSON, headers and bytes read.
|
| 68 |
+
|
| 69 |
+
``after`` runs once the response is sent, inside the same event loop.
|
| 70 |
+
"""
|
| 71 |
+
chunks = [body[i : i + CHUNK] for i in range(0, len(body), CHUNK)]
|
| 72 |
+
result: dict[str, Any] = {"read": 0, "body": b""}
|
| 73 |
+
next_chunk = 0
|
| 74 |
+
|
| 75 |
+
async def receive() -> Message:
|
| 76 |
+
nonlocal next_chunk
|
| 77 |
+
if next_chunk == len(chunks):
|
| 78 |
+
await anyio.sleep_forever() # a live connection with nothing more to send
|
| 79 |
+
chunk = chunks[next_chunk]
|
| 80 |
+
next_chunk += 1
|
| 81 |
+
result["read"] += len(chunk)
|
| 82 |
+
return {"type": "http.request", "body": chunk, "more_body": next_chunk < len(chunks)}
|
| 83 |
+
|
| 84 |
+
async def send(message: Message) -> None:
|
| 85 |
+
if message["type"] == "http.response.start":
|
| 86 |
+
result["status"] = message["status"]
|
| 87 |
+
result["headers"] = {k.decode(): v.decode() for k, v in message["headers"]}
|
| 88 |
+
elif message["type"] == "http.response.body":
|
| 89 |
+
result["body"] += message.get("body", b"")
|
| 90 |
+
|
| 91 |
+
headers = [
|
| 92 |
+
(b"content-type", b"multipart/form-data; boundary=" + BOUNDARY),
|
| 93 |
+
(b"origin", ORIGIN.encode()),
|
| 94 |
+
]
|
| 95 |
+
if declare_length:
|
| 96 |
+
headers.append((b"content-length", str(len(body)).encode()))
|
| 97 |
+
scope = {
|
| 98 |
+
"type": "http",
|
| 99 |
+
"asgi": {"version": "3.0"},
|
| 100 |
+
"http_version": "1.1",
|
| 101 |
+
"method": "POST",
|
| 102 |
+
"scheme": "http",
|
| 103 |
+
"path": "/v1/captions",
|
| 104 |
+
"raw_path": b"/v1/captions",
|
| 105 |
+
"query_string": b"",
|
| 106 |
+
"root_path": "",
|
| 107 |
+
"headers": headers,
|
| 108 |
+
"client": ("127.0.0.1", 50000),
|
| 109 |
+
"server": ("testserver", 80),
|
| 110 |
+
}
|
| 111 |
+
|
| 112 |
+
async def run() -> None:
|
| 113 |
+
with anyio.fail_after(10):
|
| 114 |
+
await app(scope, receive, send)
|
| 115 |
+
if after is not None:
|
| 116 |
+
after()
|
| 117 |
+
|
| 118 |
+
anyio.run(run)
|
| 119 |
+
result["json"] = json.loads(result["body"])
|
| 120 |
+
return result
|
| 121 |
+
|
| 122 |
+
|
| 123 |
+
def test_declared_oversize_body_is_refused_before_it_is_read() -> None:
|
| 124 |
+
service = FakePredictorService()
|
| 125 |
+
body = _multipart(b"x" * (service.max_upload_bytes + 2 * MULTIPART_OVERHEAD_BYTES))
|
| 126 |
+
|
| 127 |
+
result = _post(_build_app(service), body, declare_length=True)
|
| 128 |
+
|
| 129 |
+
assert result["status"] == 413
|
| 130 |
+
assert result["json"] == {
|
| 131 |
+
"detail": f"Image exceeds the {service.max_upload_bytes}-byte upload limit."
|
| 132 |
+
}
|
| 133 |
+
assert result["headers"].get(REQUEST_ID_HEADER)
|
| 134 |
+
assert result["headers"].get("access-control-allow-origin") == ORIGIN
|
| 135 |
+
assert result["read"] == 0
|
| 136 |
+
assert service.calls == []
|
| 137 |
+
|
| 138 |
+
|
| 139 |
+
def test_undeclared_oversize_body_stops_at_the_cap() -> None:
|
| 140 |
+
service = FakePredictorService()
|
| 141 |
+
cap = service.max_upload_bytes + MULTIPART_OVERHEAD_BYTES
|
| 142 |
+
body = _multipart(b"x" * (cap * 4))
|
| 143 |
+
|
| 144 |
+
result = _post(_build_app(service), body, declare_length=False)
|
| 145 |
+
|
| 146 |
+
assert result["status"] == 413
|
| 147 |
+
assert "limit" in result["json"]["detail"]
|
| 148 |
+
assert result["read"] <= cap + CHUNK
|
| 149 |
+
assert service.calls == []
|
| 150 |
+
|
| 151 |
+
|
| 152 |
+
@pytest.mark.parametrize(("extra_bytes", "status"), [(0, 200), (1, 413)])
|
| 153 |
+
def test_body_under_the_cap_reaches_the_route(extra_bytes: int, status: int) -> None:
|
| 154 |
+
# Under the cap the middleware is transparent: the route's exact limit decides.
|
| 155 |
+
service = FakePredictorService()
|
| 156 |
+
image = b"x" * (service.max_upload_bytes + extra_bytes)
|
| 157 |
+
body = _multipart(image)
|
| 158 |
+
|
| 159 |
+
result = _post(_build_app(service), body, declare_length=True)
|
| 160 |
+
|
| 161 |
+
assert result["status"] == status
|
| 162 |
+
assert result["read"] == len(body)
|
| 163 |
+
if status == 200:
|
| 164 |
+
assert service.calls == [image]
|
| 165 |
+
else:
|
| 166 |
+
# The route's 413 and the middleware's 413 read the same.
|
| 167 |
+
assert result["json"] == {
|
| 168 |
+
"detail": f"Image exceeds the {service.max_upload_bytes}-byte upload limit."
|
| 169 |
+
}
|
| 170 |
+
assert service.calls == []
|
| 171 |
+
|
| 172 |
+
|
| 173 |
+
def test_refused_upload_leaves_no_temp_file_open(monkeypatch: pytest.MonkeyPatch) -> None:
|
| 174 |
+
# A file part past 1 MiB is spooled to disk. Starlette doesn't close it when the cap
|
| 175 |
+
# interrupts the parse, so it must be released as soon as the request ends. anyio
|
| 176 |
+
# before 4.14.2 kept it referenced from an idle worker thread.
|
| 177 |
+
spooled: weakref.WeakSet[tempfile.SpooledTemporaryFile[bytes]] = weakref.WeakSet()
|
| 178 |
+
rollovers: list[int] = []
|
| 179 |
+
original_init = tempfile.SpooledTemporaryFile.__init__
|
| 180 |
+
original_rollover = tempfile.SpooledTemporaryFile.rollover
|
| 181 |
+
|
| 182 |
+
def tracking_init(
|
| 183 |
+
self: tempfile.SpooledTemporaryFile[bytes], *args: Any, **kwargs: Any
|
| 184 |
+
) -> None:
|
| 185 |
+
original_init(self, *args, **kwargs)
|
| 186 |
+
spooled.add(self)
|
| 187 |
+
|
| 188 |
+
def tracking_rollover(self: tempfile.SpooledTemporaryFile[bytes]) -> None:
|
| 189 |
+
rollovers.append(1)
|
| 190 |
+
original_rollover(self)
|
| 191 |
+
|
| 192 |
+
monkeypatch.setattr(tempfile.SpooledTemporaryFile, "__init__", tracking_init)
|
| 193 |
+
monkeypatch.setattr(tempfile.SpooledTemporaryFile, "rollover", tracking_rollover)
|
| 194 |
+
service = FakePredictorService(max_upload_bytes=2 * 1024 * 1024)
|
| 195 |
+
body = _multipart(b"x" * (service.max_upload_bytes * 2))
|
| 196 |
+
still_open: list[int] = []
|
| 197 |
+
|
| 198 |
+
result = _post(
|
| 199 |
+
_build_app(service),
|
| 200 |
+
body,
|
| 201 |
+
declare_length=False,
|
| 202 |
+
after=lambda: still_open.append(sum(not f.closed for f in spooled)),
|
| 203 |
+
)
|
| 204 |
+
|
| 205 |
+
assert result["status"] == 413
|
| 206 |
+
assert rollovers, "the file part never reached disk, so this test proves nothing"
|
| 207 |
+
assert still_open == [0]
|
backend/app/tests/test_captions.py
ADDED
|
@@ -0,0 +1,231 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Tests for ``POST /v1/captions``.
|
| 2 |
+
|
| 3 |
+
Covers the route's status-code contract end-to-end:
|
| 4 |
+
|
| 5 |
+
* 200 — happy path, typed ``CaptionResponse`` body
|
| 6 |
+
* 400 — empty file upload
|
| 7 |
+
* 413 — payload above ``max_upload_bytes``, or image dimensions above the limits
|
| 8 |
+
* 415 — disallowed content type
|
| 9 |
+
* 422 — bytes that the predictor cannot decode (synthetic, or no readable image header)
|
| 10 |
+
* 503 — predictor not yet loaded
|
| 11 |
+
"""
|
| 12 |
+
|
| 13 |
+
from __future__ import annotations
|
| 14 |
+
|
| 15 |
+
import pytest
|
| 16 |
+
from fastapi.testclient import TestClient
|
| 17 |
+
from starlette.datastructures import UploadFile
|
| 18 |
+
|
| 19 |
+
from app.tests.conftest import FakePredictorService
|
| 20 |
+
from app.tests.image_bytes import encode, png_header
|
| 21 |
+
|
| 22 |
+
|
| 23 |
+
def _image_field(payload: bytes, content_type: str = "image/jpeg", name: str = "a.jpg"):
|
| 24 |
+
return {"image": (name, payload, content_type)}
|
| 25 |
+
|
| 26 |
+
|
| 27 |
+
def test_captions_happy_path_returns_typed_response(
|
| 28 |
+
client: TestClient, fake_service: FakePredictorService
|
| 29 |
+
) -> None:
|
| 30 |
+
response = client.post("/v1/captions", files=_image_field(b"\xff\xd8stub"))
|
| 31 |
+
assert response.status_code == 200
|
| 32 |
+
|
| 33 |
+
body = response.json()
|
| 34 |
+
assert body["caption"] == "a test caption"
|
| 35 |
+
assert body["model_version"] == "test-v0"
|
| 36 |
+
assert body["decode_strategy"] == "greedy"
|
| 37 |
+
assert body["latency_ms"] == 1.23
|
| 38 |
+
assert body["request_id"]
|
| 39 |
+
|
| 40 |
+
# Service actually received the upload payload.
|
| 41 |
+
assert fake_service.calls == [b"\xff\xd8stub"]
|
| 42 |
+
|
| 43 |
+
|
| 44 |
+
def test_captions_request_id_matches_response_header(client: TestClient) -> None:
|
| 45 |
+
response = client.post(
|
| 46 |
+
"/v1/captions",
|
| 47 |
+
files=_image_field(b"\xff\xd8stub"),
|
| 48 |
+
headers={"x-request-id": "trace-123"},
|
| 49 |
+
)
|
| 50 |
+
assert response.status_code == 200
|
| 51 |
+
assert response.headers.get("x-request-id") == "trace-123"
|
| 52 |
+
assert response.json()["request_id"] == "trace-123"
|
| 53 |
+
|
| 54 |
+
|
| 55 |
+
def test_captions_rejects_unsupported_content_type(client: TestClient) -> None:
|
| 56 |
+
response = client.post(
|
| 57 |
+
"/v1/captions",
|
| 58 |
+
files=_image_field(b"hello", content_type="text/plain", name="a.txt"),
|
| 59 |
+
)
|
| 60 |
+
assert response.status_code == 415
|
| 61 |
+
assert "Unsupported content type" in response.json()["detail"]
|
| 62 |
+
|
| 63 |
+
|
| 64 |
+
def test_captions_rejects_webp_before_reading_it(
|
| 65 |
+
client: TestClient, fake_service: FakePredictorService
|
| 66 |
+
) -> None:
|
| 67 |
+
"""TF 2.15 can't decode WebP, so it's refused up front instead of failing as a 422."""
|
| 68 |
+
response = client.post(
|
| 69 |
+
"/v1/captions",
|
| 70 |
+
files=_image_field(
|
| 71 |
+
b"RIFF\x1a\x00\x00\x00WEBPVP8 ", content_type="image/webp", name="a.webp"
|
| 72 |
+
),
|
| 73 |
+
)
|
| 74 |
+
assert response.status_code == 415
|
| 75 |
+
assert "Unsupported content type" in response.json()["detail"]
|
| 76 |
+
assert fake_service.calls == []
|
| 77 |
+
|
| 78 |
+
|
| 79 |
+
def test_captions_without_a_strategy_field_still_decodes_greedily(
|
| 80 |
+
client: TestClient, fake_service: FakePredictorService
|
| 81 |
+
) -> None:
|
| 82 |
+
"""Clients that don't send the field get exactly the old behaviour (TASK-034)."""
|
| 83 |
+
response = client.post("/v1/captions", files=_image_field(b"\xff\xd8stub"))
|
| 84 |
+
assert response.status_code == 200
|
| 85 |
+
assert response.json()["decode_strategy"] == "greedy"
|
| 86 |
+
assert fake_service.strategies == [None]
|
| 87 |
+
|
| 88 |
+
|
| 89 |
+
@pytest.mark.parametrize("strategy", ["greedy", "beam"])
|
| 90 |
+
def test_captions_passes_the_requested_strategy_through(
|
| 91 |
+
client: TestClient, fake_service: FakePredictorService, strategy: str
|
| 92 |
+
) -> None:
|
| 93 |
+
response = client.post(
|
| 94 |
+
"/v1/captions",
|
| 95 |
+
files=_image_field(b"\xff\xd8stub"),
|
| 96 |
+
data={"decode_strategy": strategy},
|
| 97 |
+
)
|
| 98 |
+
assert response.status_code == 200
|
| 99 |
+
assert response.json()["decode_strategy"] == strategy
|
| 100 |
+
assert fake_service.strategies == [strategy]
|
| 101 |
+
assert fake_service.calls == [b"\xff\xd8stub"]
|
| 102 |
+
|
| 103 |
+
|
| 104 |
+
def test_captions_treats_an_empty_strategy_field_as_omitted(
|
| 105 |
+
client: TestClient, fake_service: FakePredictorService
|
| 106 |
+
) -> None:
|
| 107 |
+
"""An empty form value is what a blank HTML control sends; it means the default."""
|
| 108 |
+
response = client.post(
|
| 109 |
+
"/v1/captions", files=_image_field(b"\xff\xd8stub"), data={"decode_strategy": ""}
|
| 110 |
+
)
|
| 111 |
+
assert response.status_code == 200
|
| 112 |
+
assert response.json()["decode_strategy"] == "greedy"
|
| 113 |
+
assert fake_service.strategies == [None]
|
| 114 |
+
|
| 115 |
+
|
| 116 |
+
@pytest.mark.parametrize("strategy", ["sampling", "BEAM", "greedy "])
|
| 117 |
+
def test_captions_rejects_an_unknown_strategy_before_reading_the_image(
|
| 118 |
+
client: TestClient, fake_service: FakePredictorService, strategy: str
|
| 119 |
+
) -> None:
|
| 120 |
+
response = client.post(
|
| 121 |
+
"/v1/captions",
|
| 122 |
+
files=_image_field(b"\xff\xd8stub"),
|
| 123 |
+
data={"decode_strategy": strategy},
|
| 124 |
+
)
|
| 125 |
+
assert response.status_code == 422
|
| 126 |
+
detail = response.json()["detail"]
|
| 127 |
+
assert detail.startswith(f"Unsupported decode_strategy: {strategy!r}.")
|
| 128 |
+
assert "['greedy', 'beam']" in detail
|
| 129 |
+
assert fake_service.calls == []
|
| 130 |
+
|
| 131 |
+
|
| 132 |
+
def test_captions_rejects_empty_upload(client: TestClient) -> None:
|
| 133 |
+
response = client.post("/v1/captions", files=_image_field(b""))
|
| 134 |
+
assert response.status_code == 400
|
| 135 |
+
assert "Empty" in response.json()["detail"]
|
| 136 |
+
|
| 137 |
+
|
| 138 |
+
def test_captions_rejects_oversize_upload(client: TestClient) -> None:
|
| 139 |
+
# Fake service.max_upload_bytes = 1024
|
| 140 |
+
response = client.post("/v1/captions", files=_image_field(b"x" * 2048))
|
| 141 |
+
assert response.status_code == 413
|
| 142 |
+
assert "limit" in response.json()["detail"].lower()
|
| 143 |
+
|
| 144 |
+
|
| 145 |
+
def test_captions_oversize_upload_reads_at_most_limit_plus_one(
|
| 146 |
+
client: TestClient,
|
| 147 |
+
fake_service: FakePredictorService,
|
| 148 |
+
monkeypatch: pytest.MonkeyPatch,
|
| 149 |
+
) -> None:
|
| 150 |
+
# TASK-006: an oversized upload must not be pulled into memory in full.
|
| 151 |
+
limit = fake_service.max_upload_bytes
|
| 152 |
+
reads: list[tuple[int, int]] = [] # (requested size, bytes returned)
|
| 153 |
+
original_read = UploadFile.read
|
| 154 |
+
|
| 155 |
+
async def spy_read(self: UploadFile, size: int = -1) -> bytes:
|
| 156 |
+
data = await original_read(self, size)
|
| 157 |
+
reads.append((size, len(data)))
|
| 158 |
+
return data
|
| 159 |
+
|
| 160 |
+
monkeypatch.setattr(UploadFile, "read", spy_read)
|
| 161 |
+
|
| 162 |
+
response = client.post("/v1/captions", files=_image_field(b"x" * (limit * 10)))
|
| 163 |
+
|
| 164 |
+
assert response.status_code == 413
|
| 165 |
+
assert reads, "route never read the upload"
|
| 166 |
+
assert all(size >= 0 for size, _ in reads), f"unbounded read() call: {reads}"
|
| 167 |
+
assert sum(n for _, n in reads) <= limit + 1
|
| 168 |
+
assert fake_service.calls == []
|
| 169 |
+
|
| 170 |
+
|
| 171 |
+
def test_captions_accepts_upload_exactly_at_limit(
|
| 172 |
+
client: TestClient, fake_service: FakePredictorService
|
| 173 |
+
) -> None:
|
| 174 |
+
payload = b"x" * fake_service.max_upload_bytes
|
| 175 |
+
response = client.post("/v1/captions", files=_image_field(payload))
|
| 176 |
+
assert response.status_code == 200
|
| 177 |
+
assert fake_service.calls == [payload]
|
| 178 |
+
|
| 179 |
+
|
| 180 |
+
def test_captions_returns_422_on_decode_failure(build_client) -> None:
|
| 181 |
+
bad_service = FakePredictorService(raise_decode_error=True)
|
| 182 |
+
with build_client(bad_service) as test_client:
|
| 183 |
+
response = test_client.post("/v1/captions", files=_image_field(b"\xff\xd8junk"))
|
| 184 |
+
assert response.status_code == 422
|
| 185 |
+
assert "synthetic decode failure" in response.json()["detail"]
|
| 186 |
+
|
| 187 |
+
|
| 188 |
+
@pytest.mark.parametrize(("fmt", "content_type"), [("JPEG", "image/jpeg"), ("PNG", "image/png")])
|
| 189 |
+
def test_captions_accepts_a_real_image_within_the_dimension_limits(
|
| 190 |
+
build_client, fmt: str, content_type: str
|
| 191 |
+
) -> None:
|
| 192 |
+
service = FakePredictorService(check_dimensions=True, max_upload_bytes=64 * 1024)
|
| 193 |
+
payload = encode(fmt)
|
| 194 |
+
with build_client(service) as test_client:
|
| 195 |
+
response = test_client.post("/v1/captions", files=_image_field(payload, content_type))
|
| 196 |
+
assert response.status_code == 200
|
| 197 |
+
assert service.calls == [payload]
|
| 198 |
+
|
| 199 |
+
|
| 200 |
+
def test_captions_returns_413_when_image_dimensions_exceed_the_limits(build_client) -> None:
|
| 201 |
+
# A 10000x5000 header in under 100 bytes: far below the byte limit.
|
| 202 |
+
service = FakePredictorService(check_dimensions=True)
|
| 203 |
+
payload = png_header(10_000, 5_000)
|
| 204 |
+
assert len(payload) < service.max_upload_bytes
|
| 205 |
+
with build_client(service) as test_client:
|
| 206 |
+
response = test_client.post(
|
| 207 |
+
"/v1/captions", files=_image_field(payload, "image/png", "big.png")
|
| 208 |
+
)
|
| 209 |
+
assert response.status_code == 413
|
| 210 |
+
assert response.json()["detail"] == (
|
| 211 |
+
"Image dimensions exceed the limit of 16384 pixels per side "
|
| 212 |
+
"and 40,000,000 pixels in total."
|
| 213 |
+
)
|
| 214 |
+
|
| 215 |
+
|
| 216 |
+
def test_captions_returns_422_for_bytes_without_an_image_header(build_client) -> None:
|
| 217 |
+
service = FakePredictorService(check_dimensions=True)
|
| 218 |
+
with build_client(service) as test_client:
|
| 219 |
+
response = test_client.post("/v1/captions", files=_image_field(b"\xff\xd8junk"))
|
| 220 |
+
assert response.status_code == 422
|
| 221 |
+
assert response.json()["detail"] == (
|
| 222 |
+
"Could not decode image bytes: expected a JPEG, PNG or BMP image."
|
| 223 |
+
)
|
| 224 |
+
|
| 225 |
+
|
| 226 |
+
def test_captions_returns_503_when_predictor_not_loaded(
|
| 227 |
+
client_without_service: TestClient,
|
| 228 |
+
) -> None:
|
| 229 |
+
response = client_without_service.post("/v1/captions", files=_image_field(b"\xff\xd8stub"))
|
| 230 |
+
assert response.status_code == 503
|
| 231 |
+
assert "not ready" in response.json()["detail"].lower()
|
backend/app/tests/test_health.py
ADDED
|
@@ -0,0 +1,56 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Tests for ``GET /healthz``.
|
| 2 |
+
|
| 3 |
+
The route reports liveness + readiness in the response body and always
|
| 4 |
+
returns 200; readiness is conveyed by ``model_loaded``.
|
| 5 |
+
"""
|
| 6 |
+
|
| 7 |
+
from __future__ import annotations
|
| 8 |
+
|
| 9 |
+
from fastapi.testclient import TestClient
|
| 10 |
+
|
| 11 |
+
from app.tests.conftest import FakePredictorService
|
| 12 |
+
|
| 13 |
+
|
| 14 |
+
def test_healthz_reports_ready_when_service_present(client: TestClient) -> None:
|
| 15 |
+
response = client.get("/healthz")
|
| 16 |
+
assert response.status_code == 200
|
| 17 |
+
|
| 18 |
+
body = response.json()
|
| 19 |
+
assert body["status"] == "ok"
|
| 20 |
+
assert body["model_loaded"] is True
|
| 21 |
+
assert body["model_version"] == "test-v0"
|
| 22 |
+
assert body["api_version"]
|
| 23 |
+
assert "timestamp" in body
|
| 24 |
+
|
| 25 |
+
|
| 26 |
+
def test_healthz_reports_loading_when_service_missing(
|
| 27 |
+
client_without_service: TestClient,
|
| 28 |
+
) -> None:
|
| 29 |
+
response = client_without_service.get("/healthz")
|
| 30 |
+
assert response.status_code == 200
|
| 31 |
+
|
| 32 |
+
body = response.json()
|
| 33 |
+
assert body["status"] == "loading"
|
| 34 |
+
assert body["model_loaded"] is False
|
| 35 |
+
|
| 36 |
+
|
| 37 |
+
def test_healthz_echoes_request_id_header(client: TestClient) -> None:
|
| 38 |
+
response = client.get("/healthz", headers={"x-request-id": "deadbeef"})
|
| 39 |
+
assert response.status_code == 200
|
| 40 |
+
assert response.headers.get("x-request-id") == "deadbeef"
|
| 41 |
+
|
| 42 |
+
|
| 43 |
+
def test_healthz_generates_request_id_when_absent(client: TestClient) -> None:
|
| 44 |
+
response = client.get("/healthz")
|
| 45 |
+
assert response.status_code == 200
|
| 46 |
+
rid = response.headers.get("x-request-id")
|
| 47 |
+
assert rid and len(rid) >= 16
|
| 48 |
+
|
| 49 |
+
|
| 50 |
+
def test_healthz_uses_overridden_model_version(
|
| 51 |
+
build_client,
|
| 52 |
+
) -> None:
|
| 53 |
+
service = FakePredictorService(model_version="v9.9.9")
|
| 54 |
+
with build_client(service) as test_client:
|
| 55 |
+
body = test_client.get("/healthz").json()
|
| 56 |
+
assert body["model_version"] == "v9.9.9"
|
backend/app/tests/test_image_guard.py
ADDED
|
@@ -0,0 +1,158 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Tests for the upload dimension guard in ``app.utils.image``.
|
| 2 |
+
|
| 3 |
+
``check_image_dimensions`` reads the image header and must reject an image
|
| 4 |
+
over the limits before ``bytes_to_tensor`` decodes it. Oversized inputs are
|
| 5 |
+
header-only bytes (``image_bytes``), so nothing large is allocated, and
|
| 6 |
+
TensorFlow is never imported here.
|
| 7 |
+
"""
|
| 8 |
+
|
| 9 |
+
from __future__ import annotations
|
| 10 |
+
|
| 11 |
+
import struct
|
| 12 |
+
import sys
|
| 13 |
+
|
| 14 |
+
import pytest
|
| 15 |
+
|
| 16 |
+
from app.tests.image_bytes import encode, jpeg_header, png_header
|
| 17 |
+
from app.utils import image as image_module
|
| 18 |
+
from app.utils.image import (
|
| 19 |
+
MAX_IMAGE_PIXELS,
|
| 20 |
+
MAX_IMAGE_SIDE,
|
| 21 |
+
ImageDecodeError,
|
| 22 |
+
ImageTooLargeError,
|
| 23 |
+
bytes_to_tensor,
|
| 24 |
+
check_image_dimensions,
|
| 25 |
+
)
|
| 26 |
+
|
| 27 |
+
LIMIT_MESSAGE = (
|
| 28 |
+
"Image dimensions exceed the limit of 16384 pixels per side and 40,000,000 pixels in total."
|
| 29 |
+
)
|
| 30 |
+
|
| 31 |
+
|
| 32 |
+
def _bmp_with_width(width: int) -> bytes:
|
| 33 |
+
data = bytearray(encode("BMP"))
|
| 34 |
+
data[18:22] = struct.pack("<i", width)
|
| 35 |
+
return bytes(data)
|
| 36 |
+
|
| 37 |
+
|
| 38 |
+
def _png_with_bad_ihdr_crc() -> bytes:
|
| 39 |
+
data = bytearray(png_header(32, 24))
|
| 40 |
+
data[29] ^= 0xFF # first byte of the IHDR checksum
|
| 41 |
+
return bytes(data)
|
| 42 |
+
|
| 43 |
+
|
| 44 |
+
@pytest.mark.parametrize("fmt", ["JPEG", "PNG", "BMP"])
|
| 45 |
+
def test_accepts_a_small_image(fmt: str) -> None:
|
| 46 |
+
assert check_image_dimensions(encode(fmt, size=(32, 24))) == (32, 24)
|
| 47 |
+
|
| 48 |
+
|
| 49 |
+
@pytest.mark.parametrize("header", [png_header, jpeg_header])
|
| 50 |
+
def test_accepts_dimensions_exactly_at_the_limits(header) -> None:
|
| 51 |
+
assert MAX_IMAGE_PIXELS == 8_000 * 5_000
|
| 52 |
+
assert check_image_dimensions(header(8_000, 5_000)) == (8_000, 5_000)
|
| 53 |
+
assert check_image_dimensions(header(MAX_IMAGE_SIDE, 1)) == (MAX_IMAGE_SIDE, 1)
|
| 54 |
+
assert check_image_dimensions(header(1, MAX_IMAGE_SIDE)) == (1, MAX_IMAGE_SIDE)
|
| 55 |
+
|
| 56 |
+
|
| 57 |
+
@pytest.mark.parametrize("header", [png_header, jpeg_header])
|
| 58 |
+
def test_rejects_a_pixel_count_over_the_limit(header) -> None:
|
| 59 |
+
# Both sides are within MAX_IMAGE_SIDE; only the product is over.
|
| 60 |
+
with pytest.raises(ImageTooLargeError):
|
| 61 |
+
check_image_dimensions(header(8_000, 5_001))
|
| 62 |
+
|
| 63 |
+
|
| 64 |
+
@pytest.mark.parametrize("header", [png_header, jpeg_header])
|
| 65 |
+
@pytest.mark.parametrize("size", [(MAX_IMAGE_SIDE + 1, 1), (1, MAX_IMAGE_SIDE + 1)])
|
| 66 |
+
def test_rejects_a_side_over_the_limit(header, size: tuple[int, int]) -> None:
|
| 67 |
+
with pytest.raises(ImageTooLargeError):
|
| 68 |
+
check_image_dimensions(header(*size))
|
| 69 |
+
|
| 70 |
+
|
| 71 |
+
def test_rejects_dimensions_past_pillows_own_bomb_ceiling() -> None:
|
| 72 |
+
# 65535 x 65535 makes Image.open itself raise DecompressionBombError.
|
| 73 |
+
with pytest.raises(ImageTooLargeError):
|
| 74 |
+
check_image_dimensions(png_header(65_535, 65_535))
|
| 75 |
+
|
| 76 |
+
|
| 77 |
+
def test_limit_message_is_concise_and_names_only_the_limits() -> None:
|
| 78 |
+
with pytest.raises(ImageTooLargeError) as info:
|
| 79 |
+
check_image_dimensions(png_header(10_000, 5_000))
|
| 80 |
+
assert str(info.value) == LIMIT_MESSAGE
|
| 81 |
+
|
| 82 |
+
|
| 83 |
+
def test_honours_explicit_limits() -> None:
|
| 84 |
+
small = encode("PNG", size=(32, 24))
|
| 85 |
+
with pytest.raises(ImageTooLargeError):
|
| 86 |
+
check_image_dimensions(small, max_side=31)
|
| 87 |
+
with pytest.raises(ImageTooLargeError):
|
| 88 |
+
check_image_dimensions(small, max_pixels=32 * 24 - 1)
|
| 89 |
+
assert check_image_dimensions(small, max_side=32, max_pixels=32 * 24) == (32, 24)
|
| 90 |
+
|
| 91 |
+
|
| 92 |
+
@pytest.mark.parametrize(
|
| 93 |
+
("data", "error"),
|
| 94 |
+
[
|
| 95 |
+
pytest.param(b"", ImageDecodeError, id="empty"),
|
| 96 |
+
pytest.param(b"not an image at all", ImageDecodeError, id="garbage"),
|
| 97 |
+
pytest.param(b"\xff\xd8\xff", ImageDecodeError, id="jpeg-soi-only"),
|
| 98 |
+
pytest.param(jpeg_header(32, 24)[:10], ImageDecodeError, id="jpeg-cut-in-sof"),
|
| 99 |
+
pytest.param(png_header(32, 24)[:20], ImageDecodeError, id="png-cut-in-ihdr"),
|
| 100 |
+
pytest.param(png_header(32, 24)[:33], ImageDecodeError, id="png-ends-after-ihdr"),
|
| 101 |
+
pytest.param(_png_with_bad_ihdr_crc(), ImageDecodeError, id="png-bad-ihdr-crc"),
|
| 102 |
+
pytest.param(png_header(0, 24), ImageDecodeError, id="png-zero-width"),
|
| 103 |
+
pytest.param(png_header(32, 0), ImageDecodeError, id="png-zero-height"),
|
| 104 |
+
pytest.param(jpeg_header(0, 24), ImageDecodeError, id="jpeg-zero-width"),
|
| 105 |
+
pytest.param(jpeg_header(32, 0), ImageDecodeError, id="jpeg-zero-height"),
|
| 106 |
+
# Pillow reads a BMP width as unsigned, so -32 claims about 4.3e9 pixels.
|
| 107 |
+
pytest.param(_bmp_with_width(-32), ImageTooLargeError, id="bmp-negative-width"),
|
| 108 |
+
],
|
| 109 |
+
)
|
| 110 |
+
def test_malformed_or_invalid_headers_fail_with_a_client_error(data: bytes, error: type) -> None:
|
| 111 |
+
with pytest.raises(ImageDecodeError) as info:
|
| 112 |
+
check_image_dimensions(data)
|
| 113 |
+
assert type(info.value) is error
|
| 114 |
+
|
| 115 |
+
|
| 116 |
+
@pytest.mark.parametrize("size", [(0, 24), (32, 0), (-1, 24)])
|
| 117 |
+
def test_rejects_non_positive_dimensions_from_the_parser(
|
| 118 |
+
monkeypatch: pytest.MonkeyPatch, size: tuple[int, int]
|
| 119 |
+
) -> None:
|
| 120 |
+
# Pillow already refuses these headers; the guard checks again on its own.
|
| 121 |
+
class _Parsed:
|
| 122 |
+
def __init__(self) -> None:
|
| 123 |
+
self.size = size
|
| 124 |
+
|
| 125 |
+
def __enter__(self) -> _Parsed:
|
| 126 |
+
return self
|
| 127 |
+
|
| 128 |
+
def __exit__(self, *exc_info: object) -> None:
|
| 129 |
+
return None
|
| 130 |
+
|
| 131 |
+
monkeypatch.setattr(image_module.Image, "open", lambda *args, **kwargs: _Parsed())
|
| 132 |
+
with pytest.raises(ImageDecodeError) as info:
|
| 133 |
+
check_image_dimensions(b"stub")
|
| 134 |
+
assert type(info.value) is ImageDecodeError
|
| 135 |
+
|
| 136 |
+
|
| 137 |
+
@pytest.mark.parametrize("fmt", ["GIF", "WEBP", "TIFF"])
|
| 138 |
+
def test_rejects_formats_outside_the_header_check(fmt: str) -> None:
|
| 139 |
+
# A GIF is decoded at its largest frame's size, which its header doesn't
|
| 140 |
+
# give, and opening a WebP allocates its canvas, so neither is parsed.
|
| 141 |
+
# TF 2.15 can't decode WebP or TIFF anyway.
|
| 142 |
+
with pytest.raises(ImageDecodeError) as info:
|
| 143 |
+
check_image_dimensions(encode(fmt))
|
| 144 |
+
assert type(info.value) is ImageDecodeError
|
| 145 |
+
|
| 146 |
+
|
| 147 |
+
def test_bytes_to_tensor_checks_dimensions_before_importing_tensorflow(
|
| 148 |
+
monkeypatch: pytest.MonkeyPatch,
|
| 149 |
+
) -> None:
|
| 150 |
+
# With TensorFlow unimportable, reaching the decoder raises ImportError.
|
| 151 |
+
monkeypatch.setitem(sys.modules, "tensorflow", None)
|
| 152 |
+
|
| 153 |
+
with pytest.raises(ImageTooLargeError):
|
| 154 |
+
bytes_to_tensor(png_header(10_000, 5_000))
|
| 155 |
+
with pytest.raises(ImageDecodeError):
|
| 156 |
+
bytes_to_tensor(b"not an image at all")
|
| 157 |
+
with pytest.raises(ImportError):
|
| 158 |
+
bytes_to_tensor(encode("PNG")) # within the limits: goes on to decode
|
backend/app/tests/test_inference_admission.py
ADDED
|
@@ -0,0 +1,325 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Tests for bounded inference admission (TASK-027).
|
| 2 |
+
|
| 3 |
+
``PredictorService`` runs with a gated stand-in predictor: its
|
| 4 |
+
``predict_tensor`` blocks a real worker thread until the test opens the gate,
|
| 5 |
+
so requests genuinely overlap. Decoding runs the real TASK-026 header guard
|
| 6 |
+
and skips TensorFlow. Nothing sleeps: tests wait for the worker threads' own
|
| 7 |
+
signals, and ``fail_after`` only guards against a hang.
|
| 8 |
+
"""
|
| 9 |
+
|
| 10 |
+
from __future__ import annotations
|
| 11 |
+
|
| 12 |
+
import math
|
| 13 |
+
import threading
|
| 14 |
+
from typing import Any, cast
|
| 15 |
+
|
| 16 |
+
import anyio
|
| 17 |
+
import httpx
|
| 18 |
+
import pytest
|
| 19 |
+
from anyio import from_thread, to_thread
|
| 20 |
+
|
| 21 |
+
from app.api.routes import BUSY_DETAIL, BUSY_RETRY_AFTER_SECONDS
|
| 22 |
+
from app.services import predictor_service
|
| 23 |
+
from app.services.predictor_service import (
|
| 24 |
+
MAX_CONCURRENT_INFERENCES,
|
| 25 |
+
PredictorBusyError,
|
| 26 |
+
PredictorService,
|
| 27 |
+
)
|
| 28 |
+
from app.tests.conftest import _build_app
|
| 29 |
+
from app.tests.image_bytes import encode, png_header
|
| 30 |
+
from app.utils.image import ImageTooLargeError, check_image_dimensions
|
| 31 |
+
|
| 32 |
+
IMAGE = encode("PNG")
|
| 33 |
+
HANG_GUARD_SECONDS = 10 # reached only if a test is broken
|
| 34 |
+
SLOTS = MAX_CONCURRENT_INFERENCES
|
| 35 |
+
|
| 36 |
+
|
| 37 |
+
class InferenceFailure(Exception):
|
| 38 |
+
"""A predictor failure that can't be mistaken for ``PredictorBusyError``."""
|
| 39 |
+
|
| 40 |
+
|
| 41 |
+
class GatedPredictor:
|
| 42 |
+
"""``predict_tensor`` holds its worker thread until ``open_gate()``."""
|
| 43 |
+
|
| 44 |
+
decode_strategy = "greedy"
|
| 45 |
+
|
| 46 |
+
def __init__(self) -> None:
|
| 47 |
+
self.error: Exception | None = None
|
| 48 |
+
self.running = 0
|
| 49 |
+
self.peak = 0
|
| 50 |
+
self.calls = 0
|
| 51 |
+
self.strategies: list[str | None] = []
|
| 52 |
+
self._gate = threading.Event()
|
| 53 |
+
self._lock = threading.Lock()
|
| 54 |
+
self._entered, self._entries = anyio.create_memory_object_stream[None](math.inf)
|
| 55 |
+
|
| 56 |
+
def predict_tensor(self, tensor: object, *, decode_strategy: str | None = None) -> str:
|
| 57 |
+
with self._lock:
|
| 58 |
+
self.calls += 1
|
| 59 |
+
self.strategies.append(decode_strategy)
|
| 60 |
+
self.running += 1
|
| 61 |
+
self.peak = max(self.peak, self.running)
|
| 62 |
+
try:
|
| 63 |
+
from_thread.run_sync(self._entered.send_nowait, None)
|
| 64 |
+
if not self._gate.wait(HANG_GUARD_SECONDS):
|
| 65 |
+
raise TimeoutError("the test never opened the gate")
|
| 66 |
+
if self.error is not None:
|
| 67 |
+
raise self.error
|
| 68 |
+
return "a gated caption"
|
| 69 |
+
finally:
|
| 70 |
+
with self._lock:
|
| 71 |
+
self.running -= 1
|
| 72 |
+
|
| 73 |
+
def open_gate(self) -> None:
|
| 74 |
+
self._gate.set()
|
| 75 |
+
|
| 76 |
+
def close_gate(self) -> None:
|
| 77 |
+
self._gate.clear()
|
| 78 |
+
|
| 79 |
+
async def wait_running(self, count: int) -> None:
|
| 80 |
+
"""Return once ``count`` more calls have entered ``predict_tensor``."""
|
| 81 |
+
for _ in range(count):
|
| 82 |
+
await self._entries.receive()
|
| 83 |
+
|
| 84 |
+
|
| 85 |
+
@pytest.fixture(autouse=True)
|
| 86 |
+
def decoded(monkeypatch: pytest.MonkeyPatch) -> list[bytes]:
|
| 87 |
+
"""Replace the TF decode with the real header guard plus a placeholder tensor."""
|
| 88 |
+
images: list[bytes] = []
|
| 89 |
+
|
| 90 |
+
def decode(image_bytes: bytes) -> str:
|
| 91 |
+
check_image_dimensions(image_bytes)
|
| 92 |
+
images.append(image_bytes)
|
| 93 |
+
return "tensor"
|
| 94 |
+
|
| 95 |
+
monkeypatch.setattr(predictor_service, "bytes_to_tensor", decode)
|
| 96 |
+
return images
|
| 97 |
+
|
| 98 |
+
|
| 99 |
+
def _service(predictor: GatedPredictor, slots: int = SLOTS) -> PredictorService:
|
| 100 |
+
return PredictorService(
|
| 101 |
+
predictor=cast(Any, predictor),
|
| 102 |
+
model_version="test-v0",
|
| 103 |
+
max_upload_bytes=64 * 1024,
|
| 104 |
+
max_concurrent_inferences=slots,
|
| 105 |
+
)
|
| 106 |
+
|
| 107 |
+
|
| 108 |
+
async def _saturate_then_release(
|
| 109 |
+
service: PredictorService, predictor: GatedPredictor, slots: int = SLOTS
|
| 110 |
+
) -> list[str]:
|
| 111 |
+
"""Fill every slot with a blocked request, check one more is refused, then finish them."""
|
| 112 |
+
captions: list[str] = []
|
| 113 |
+
|
| 114 |
+
async def caption() -> None:
|
| 115 |
+
text, _ = await service.caption_image_bytes(IMAGE)
|
| 116 |
+
captions.append(text)
|
| 117 |
+
|
| 118 |
+
with anyio.fail_after(HANG_GUARD_SECONDS):
|
| 119 |
+
async with anyio.create_task_group() as tg:
|
| 120 |
+
for _ in range(slots):
|
| 121 |
+
tg.start_soon(caption)
|
| 122 |
+
await predictor.wait_running(slots)
|
| 123 |
+
assert service.inferences_in_flight == slots
|
| 124 |
+
# Refused while every slot is still blocked, so it didn't wait for one.
|
| 125 |
+
with pytest.raises(PredictorBusyError):
|
| 126 |
+
await service.caption_image_bytes(IMAGE)
|
| 127 |
+
predictor.open_gate()
|
| 128 |
+
predictor.close_gate()
|
| 129 |
+
assert service.inferences_in_flight == 0
|
| 130 |
+
return captions
|
| 131 |
+
|
| 132 |
+
|
| 133 |
+
def test_concurrent_requests_keep_their_own_decode_strategy() -> None:
|
| 134 |
+
"""Two in-flight requests with different strategies never share a setting (TASK-034)."""
|
| 135 |
+
|
| 136 |
+
async def main() -> None:
|
| 137 |
+
predictor = GatedPredictor()
|
| 138 |
+
service = _service(predictor, slots=2)
|
| 139 |
+
results: dict[str, str] = {}
|
| 140 |
+
|
| 141 |
+
async def caption(strategy: str | None) -> None:
|
| 142 |
+
text, _ = await service.caption_image_bytes(IMAGE, decode_strategy=strategy)
|
| 143 |
+
results[str(strategy)] = text
|
| 144 |
+
|
| 145 |
+
with anyio.fail_after(HANG_GUARD_SECONDS):
|
| 146 |
+
async with anyio.create_task_group() as tg:
|
| 147 |
+
tg.start_soon(caption, "beam")
|
| 148 |
+
tg.start_soon(caption, None)
|
| 149 |
+
await predictor.wait_running(2)
|
| 150 |
+
assert predictor.running == 2
|
| 151 |
+
predictor.open_gate()
|
| 152 |
+
assert sorted(predictor.strategies, key=str) == [None, "beam"]
|
| 153 |
+
assert results == {"beam": "a gated caption", "None": "a gated caption"}
|
| 154 |
+
assert service.decode_strategy == "greedy" # the predictor's default is untouched
|
| 155 |
+
|
| 156 |
+
anyio.run(main)
|
| 157 |
+
|
| 158 |
+
|
| 159 |
+
def test_rejects_a_limit_below_one() -> None:
|
| 160 |
+
with pytest.raises(ValueError, match="max_concurrent_inferences"):
|
| 161 |
+
_service(GatedPredictor(), slots=0)
|
| 162 |
+
|
| 163 |
+
|
| 164 |
+
def test_a_request_is_admitted_when_a_slot_is_free() -> None:
|
| 165 |
+
async def main() -> None:
|
| 166 |
+
predictor = GatedPredictor()
|
| 167 |
+
predictor.open_gate()
|
| 168 |
+
service = _service(predictor)
|
| 169 |
+
with anyio.fail_after(HANG_GUARD_SECONDS):
|
| 170 |
+
caption, latency_ms = await service.caption_image_bytes(IMAGE)
|
| 171 |
+
assert caption == "a gated caption"
|
| 172 |
+
assert latency_ms >= 0
|
| 173 |
+
assert service.inferences_in_flight == 0
|
| 174 |
+
|
| 175 |
+
anyio.run(main)
|
| 176 |
+
|
| 177 |
+
|
| 178 |
+
def test_a_request_over_capacity_is_refused_without_decoding(decoded: list[bytes]) -> None:
|
| 179 |
+
async def main() -> None:
|
| 180 |
+
predictor = GatedPredictor()
|
| 181 |
+
service = _service(predictor)
|
| 182 |
+
captions = await _saturate_then_release(service, predictor)
|
| 183 |
+
assert captions == ["a gated caption"] * SLOTS
|
| 184 |
+
assert len(decoded) == SLOTS # the refused request decoded nothing
|
| 185 |
+
assert predictor.calls == SLOTS
|
| 186 |
+
assert predictor.peak == SLOTS
|
| 187 |
+
|
| 188 |
+
anyio.run(main)
|
| 189 |
+
|
| 190 |
+
|
| 191 |
+
def test_slots_are_released_after_successful_inference() -> None:
|
| 192 |
+
async def main() -> None:
|
| 193 |
+
predictor = GatedPredictor()
|
| 194 |
+
service = _service(predictor)
|
| 195 |
+
await _saturate_then_release(service, predictor)
|
| 196 |
+
# Every slot came back: a second full round is admitted the same way.
|
| 197 |
+
await _saturate_then_release(service, predictor)
|
| 198 |
+
assert predictor.calls == 2 * SLOTS
|
| 199 |
+
assert predictor.peak == SLOTS
|
| 200 |
+
|
| 201 |
+
anyio.run(main)
|
| 202 |
+
|
| 203 |
+
|
| 204 |
+
def test_slot_is_released_when_inference_raises() -> None:
|
| 205 |
+
async def main() -> None:
|
| 206 |
+
predictor = GatedPredictor()
|
| 207 |
+
predictor.error = InferenceFailure("model failed")
|
| 208 |
+
predictor.open_gate()
|
| 209 |
+
service = _service(predictor)
|
| 210 |
+
with anyio.fail_after(HANG_GUARD_SECONDS), pytest.raises(InferenceFailure):
|
| 211 |
+
await service.caption_image_bytes(IMAGE)
|
| 212 |
+
assert service.inferences_in_flight == 0
|
| 213 |
+
|
| 214 |
+
predictor.error = None
|
| 215 |
+
predictor.close_gate()
|
| 216 |
+
await _saturate_then_release(service, predictor)
|
| 217 |
+
|
| 218 |
+
anyio.run(main)
|
| 219 |
+
|
| 220 |
+
|
| 221 |
+
def test_slot_is_released_when_decoding_fails() -> None:
|
| 222 |
+
async def main() -> None:
|
| 223 |
+
predictor = GatedPredictor()
|
| 224 |
+
service = _service(predictor)
|
| 225 |
+
with anyio.fail_after(HANG_GUARD_SECONDS), pytest.raises(ImageTooLargeError):
|
| 226 |
+
await service.caption_image_bytes(png_header(10_000, 5_000))
|
| 227 |
+
assert service.inferences_in_flight == 0
|
| 228 |
+
assert predictor.calls == 0
|
| 229 |
+
|
| 230 |
+
await _saturate_then_release(service, predictor)
|
| 231 |
+
|
| 232 |
+
anyio.run(main)
|
| 233 |
+
|
| 234 |
+
|
| 235 |
+
def test_a_cancelled_request_keeps_its_slot_until_its_thread_finishes() -> None:
|
| 236 |
+
async def main() -> None:
|
| 237 |
+
predictor = GatedPredictor()
|
| 238 |
+
service = _service(predictor, slots=1)
|
| 239 |
+
outcome: list[str] = []
|
| 240 |
+
scope = anyio.CancelScope()
|
| 241 |
+
|
| 242 |
+
async def request() -> None:
|
| 243 |
+
with scope:
|
| 244 |
+
await service.caption_image_bytes(IMAGE)
|
| 245 |
+
outcome.append("ended")
|
| 246 |
+
|
| 247 |
+
with anyio.fail_after(HANG_GUARD_SECONDS):
|
| 248 |
+
async with anyio.create_task_group() as tg:
|
| 249 |
+
tg.start_soon(request)
|
| 250 |
+
await predictor.wait_running(1)
|
| 251 |
+
scope.cancel()
|
| 252 |
+
await anyio.wait_all_tasks_blocked()
|
| 253 |
+
# Cancelled, but its thread is still running, so the slot stays taken
|
| 254 |
+
# and no second inference can start beside it.
|
| 255 |
+
assert outcome == []
|
| 256 |
+
assert predictor.running == 1
|
| 257 |
+
assert service.inferences_in_flight == 1
|
| 258 |
+
with pytest.raises(PredictorBusyError):
|
| 259 |
+
await service.caption_image_bytes(IMAGE)
|
| 260 |
+
predictor.open_gate()
|
| 261 |
+
|
| 262 |
+
# It ended only once its thread had finished. The cancellation can't stop
|
| 263 |
+
# that thread, and here it never surfaced: no checkpoint followed the wait.
|
| 264 |
+
assert outcome == ["ended"]
|
| 265 |
+
assert service.inferences_in_flight == 0
|
| 266 |
+
assert predictor.running == 0
|
| 267 |
+
assert predictor.peak == 1
|
| 268 |
+
caption, _ = await service.caption_image_bytes(IMAGE)
|
| 269 |
+
assert caption == "a gated caption"
|
| 270 |
+
|
| 271 |
+
anyio.run(main)
|
| 272 |
+
|
| 273 |
+
|
| 274 |
+
def _upload(payload: bytes, content_type: str = "image/png") -> dict[str, Any]:
|
| 275 |
+
return {"image": ("a.png", payload, content_type)}
|
| 276 |
+
|
| 277 |
+
|
| 278 |
+
def test_overload_is_a_fast_503_and_healthz_needs_no_worker_thread() -> None:
|
| 279 |
+
async def main() -> None:
|
| 280 |
+
predictor = GatedPredictor()
|
| 281 |
+
service = _service(predictor)
|
| 282 |
+
# Inference will hold every worker thread anyio has: a route that needed
|
| 283 |
+
# one would hang here until the hang guard fails the test.
|
| 284 |
+
to_thread.current_default_thread_limiter().total_tokens = SLOTS
|
| 285 |
+
app = _build_app(cast(Any, service))
|
| 286 |
+
# httpx types its ASGI app more narrowly than Starlette's own signature.
|
| 287 |
+
transport = httpx.ASGITransport(app=cast(Any, app))
|
| 288 |
+
responses: list[httpx.Response] = []
|
| 289 |
+
|
| 290 |
+
async with httpx.AsyncClient(transport=transport, base_url="http://test") as client:
|
| 291 |
+
|
| 292 |
+
async def post() -> None:
|
| 293 |
+
responses.append(await client.post("/v1/captions", files=_upload(IMAGE)))
|
| 294 |
+
|
| 295 |
+
with anyio.fail_after(HANG_GUARD_SECONDS):
|
| 296 |
+
async with anyio.create_task_group() as tg:
|
| 297 |
+
for _ in range(SLOTS):
|
| 298 |
+
tg.start_soon(post)
|
| 299 |
+
await predictor.wait_running(SLOTS)
|
| 300 |
+
assert to_thread.current_default_thread_limiter().available_tokens == 0
|
| 301 |
+
|
| 302 |
+
health = await client.get("/healthz")
|
| 303 |
+
assert health.status_code == 200
|
| 304 |
+
assert health.json()["status"] == "ok"
|
| 305 |
+
assert health.json()["model_loaded"] is True
|
| 306 |
+
|
| 307 |
+
busy = await client.post("/v1/captions", files=_upload(IMAGE))
|
| 308 |
+
assert busy.status_code == 503
|
| 309 |
+
assert busy.json()["detail"] == BUSY_DETAIL
|
| 310 |
+
assert busy.headers["retry-after"] == str(BUSY_RETRY_AFTER_SECONDS)
|
| 311 |
+
assert busy.headers["x-request-id"]
|
| 312 |
+
|
| 313 |
+
# Upload validation still runs before admission.
|
| 314 |
+
wrong_type = await client.post(
|
| 315 |
+
"/v1/captions", files=_upload(b"hello", "text/plain")
|
| 316 |
+
)
|
| 317 |
+
assert wrong_type.status_code == 415
|
| 318 |
+
|
| 319 |
+
predictor.open_gate()
|
| 320 |
+
|
| 321 |
+
assert [response.status_code for response in responses] == [200] * SLOTS
|
| 322 |
+
assert all(r.json()["caption"] == "a gated caption" for r in responses)
|
| 323 |
+
assert predictor.peak == SLOTS
|
| 324 |
+
|
| 325 |
+
anyio.run(main)
|
backend/app/tests/test_startup_integrity.py
ADDED
|
@@ -0,0 +1,154 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Startup never hands an unverified file to the model or tokenizer loader (TASK-029).
|
| 2 |
+
|
| 3 |
+
These tests run the real ``lifespan`` and the real ``resolve_weights`` with its
|
| 4 |
+
default downloader. Only ``huggingface_hub.snapshot_download`` (no network) and
|
| 5 |
+
``CaptionPredictor`` (no TensorFlow) are replaced. The stand-in predictor reads
|
| 6 |
+
the files the way the real one does: the vocabulary through ``read_vocabulary``,
|
| 7 |
+
which ``CaptionTokenizer.load`` calls first, and the weights file's bytes.
|
| 8 |
+
|
| 9 |
+
Importing ``app.main`` stays TensorFlow-free only while ``captioning.inference``
|
| 10 |
+
imports TensorFlow inside functions. If that changes, move this module to
|
| 11 |
+
``tests/unit`` so the backend suite stays TensorFlow-free.
|
| 12 |
+
"""
|
| 13 |
+
|
| 14 |
+
from __future__ import annotations
|
| 15 |
+
|
| 16 |
+
import hashlib
|
| 17 |
+
import json
|
| 18 |
+
from pathlib import Path
|
| 19 |
+
from typing import Any, ClassVar
|
| 20 |
+
|
| 21 |
+
import anyio
|
| 22 |
+
import huggingface_hub
|
| 23 |
+
import pytest
|
| 24 |
+
from fastapi import FastAPI
|
| 25 |
+
|
| 26 |
+
import app.main as app_main
|
| 27 |
+
from app.core.config import BackendSettings
|
| 28 |
+
from app.services.weights_loader import PINNED_ARTIFACT_SHA256, ArtifactIntegrityError
|
| 29 |
+
from captioning.config.schema import AppConfig
|
| 30 |
+
from captioning.preprocessing.tokenizer import read_vocabulary
|
| 31 |
+
|
| 32 |
+
REPO = "user/captioning-weights"
|
| 33 |
+
COMMIT = "0123456789abcdef0123456789abcdef01234567"
|
| 34 |
+
UNPINNED_COMMIT = "fedcba9876543210fedcba9876543210fedcba98"
|
| 35 |
+
WEIGHTS = b"\x89HDF\r\n\x1a\nfake weights"
|
| 36 |
+
VOCAB = ["", "[UNK]", "a", "dog", "on", "grass"]
|
| 37 |
+
|
| 38 |
+
|
| 39 |
+
class RecordingPredictor:
|
| 40 |
+
"""Stands in for ``CaptionPredictor``; records every file it consumes."""
|
| 41 |
+
|
| 42 |
+
consumed: ClassVar[list[str]] = []
|
| 43 |
+
vocabulary: ClassVar[list[str]] = []
|
| 44 |
+
|
| 45 |
+
@classmethod
|
| 46 |
+
def from_artifacts(
|
| 47 |
+
cls, *, weights_path: Path, tokenizer_dir: Path, config: AppConfig
|
| 48 |
+
) -> RecordingPredictor:
|
| 49 |
+
cls.vocabulary = read_vocabulary(tokenizer_dir, config.model.vocabulary_size)
|
| 50 |
+
cls.consumed.append("vocab.json")
|
| 51 |
+
Path(weights_path).read_bytes()
|
| 52 |
+
cls.consumed.append(Path(weights_path).name)
|
| 53 |
+
return cls()
|
| 54 |
+
|
| 55 |
+
def warmup(self) -> None:
|
| 56 |
+
pass
|
| 57 |
+
|
| 58 |
+
|
| 59 |
+
@pytest.fixture
|
| 60 |
+
def snapshot(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path:
|
| 61 |
+
"""A snapshot holding the pinned bytes, served by a no-network downloader."""
|
| 62 |
+
directory = tmp_path / "snapshots" / COMMIT
|
| 63 |
+
directory.mkdir(parents=True)
|
| 64 |
+
vocab_bytes = json.dumps(VOCAB).encode()
|
| 65 |
+
(directory / "model.h5").write_bytes(WEIGHTS)
|
| 66 |
+
(directory / "vocab.json").write_bytes(vocab_bytes)
|
| 67 |
+
monkeypatch.setitem(
|
| 68 |
+
PINNED_ARTIFACT_SHA256,
|
| 69 |
+
(REPO, COMMIT),
|
| 70 |
+
{
|
| 71 |
+
"model.h5": hashlib.sha256(WEIGHTS).hexdigest(),
|
| 72 |
+
"vocab.json": hashlib.sha256(vocab_bytes).hexdigest(),
|
| 73 |
+
},
|
| 74 |
+
)
|
| 75 |
+
return directory
|
| 76 |
+
|
| 77 |
+
|
| 78 |
+
@pytest.fixture
|
| 79 |
+
def downloads(snapshot: Path, monkeypatch: pytest.MonkeyPatch) -> list[dict[str, Any]]:
|
| 80 |
+
calls: list[dict[str, Any]] = []
|
| 81 |
+
|
| 82 |
+
def fake_snapshot_download(**kwargs: Any) -> str:
|
| 83 |
+
calls.append(kwargs)
|
| 84 |
+
return str(snapshot)
|
| 85 |
+
|
| 86 |
+
monkeypatch.setattr(huggingface_hub, "snapshot_download", fake_snapshot_download)
|
| 87 |
+
monkeypatch.setattr(app_main, "CaptionPredictor", RecordingPredictor)
|
| 88 |
+
monkeypatch.setattr(RecordingPredictor, "consumed", [])
|
| 89 |
+
monkeypatch.setattr(RecordingPredictor, "vocabulary", [])
|
| 90 |
+
return calls
|
| 91 |
+
|
| 92 |
+
|
| 93 |
+
def _start(revision: str = COMMIT) -> FastAPI:
|
| 94 |
+
"""Run the lifespan's startup and shutdown; return the app as it was while serving."""
|
| 95 |
+
app = FastAPI()
|
| 96 |
+
app.state.backend_settings = BackendSettings(
|
| 97 |
+
weights_hub_repo=REPO, weights_hub_revision=revision, warmup=True
|
| 98 |
+
)
|
| 99 |
+
app.state.app_config = AppConfig()
|
| 100 |
+
app.state.predictor_service = None
|
| 101 |
+
seen: dict[str, Any] = {}
|
| 102 |
+
|
| 103 |
+
async def run() -> None:
|
| 104 |
+
async with app_main.lifespan(app):
|
| 105 |
+
seen["service"] = app.state.predictor_service
|
| 106 |
+
|
| 107 |
+
anyio.run(run)
|
| 108 |
+
app.state.seen_service = seen["service"]
|
| 109 |
+
return app
|
| 110 |
+
|
| 111 |
+
|
| 112 |
+
def test_verified_files_load(downloads: list[dict[str, Any]]) -> None:
|
| 113 |
+
app = _start()
|
| 114 |
+
|
| 115 |
+
assert app.state.seen_service is not None
|
| 116 |
+
assert RecordingPredictor.vocabulary == VOCAB
|
| 117 |
+
assert RecordingPredictor.consumed == ["vocab.json", "model.h5"]
|
| 118 |
+
assert [call["allow_patterns"] for call in downloads] == [["model.h5", "vocab.json"]]
|
| 119 |
+
assert [call["revision"] for call in downloads] == [COMMIT]
|
| 120 |
+
|
| 121 |
+
|
| 122 |
+
def test_a_tampered_weights_file_stops_startup_before_the_model_loads(
|
| 123 |
+
snapshot: Path, downloads: list[dict[str, Any]]
|
| 124 |
+
) -> None:
|
| 125 |
+
(snapshot / "model.h5").write_bytes(WEIGHTS + b"\x00")
|
| 126 |
+
|
| 127 |
+
with pytest.raises(ArtifactIntegrityError, match="model.h5 has SHA-256"):
|
| 128 |
+
_start()
|
| 129 |
+
|
| 130 |
+
assert RecordingPredictor.consumed == []
|
| 131 |
+
|
| 132 |
+
|
| 133 |
+
def test_a_tampered_vocabulary_is_never_read_even_when_it_would_validate(
|
| 134 |
+
snapshot: Path, downloads: list[dict[str, Any]]
|
| 135 |
+
) -> None:
|
| 136 |
+
"""Swapping two tokens keeps the list valid but shifts token ids."""
|
| 137 |
+
swapped = VOCAB[:2] + [VOCAB[3], VOCAB[2]] + VOCAB[4:]
|
| 138 |
+
(snapshot / "vocab.json").write_bytes(json.dumps(swapped).encode())
|
| 139 |
+
assert read_vocabulary(snapshot, AppConfig().model.vocabulary_size) == swapped
|
| 140 |
+
|
| 141 |
+
with pytest.raises(ArtifactIntegrityError, match="vocab.json has SHA-256"):
|
| 142 |
+
_start()
|
| 143 |
+
|
| 144 |
+
assert RecordingPredictor.consumed == []
|
| 145 |
+
|
| 146 |
+
|
| 147 |
+
def test_an_unpinned_commit_stops_startup_before_any_download(
|
| 148 |
+
downloads: list[dict[str, Any]],
|
| 149 |
+
) -> None:
|
| 150 |
+
with pytest.raises(ValueError, match="has no SHA-256 digests"):
|
| 151 |
+
_start(UNPINNED_COMMIT)
|
| 152 |
+
|
| 153 |
+
assert downloads == []
|
| 154 |
+
assert RecordingPredictor.consumed == []
|
backend/app/tests/test_weights_loader.py
ADDED
|
@@ -0,0 +1,400 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Unit tests for ``app.services.weights_loader.resolve_weights``.
|
| 2 |
+
|
| 3 |
+
These tests never hit the network — the downloader is injected as a stub.
|
| 4 |
+
"""
|
| 5 |
+
|
| 6 |
+
from __future__ import annotations
|
| 7 |
+
|
| 8 |
+
import hashlib
|
| 9 |
+
import json
|
| 10 |
+
from pathlib import Path
|
| 11 |
+
|
| 12 |
+
import pytest
|
| 13 |
+
|
| 14 |
+
from app.core.config import BackendSettings
|
| 15 |
+
from app.services.weights_loader import (
|
| 16 |
+
PINNED_ARTIFACT_SHA256,
|
| 17 |
+
PINNED_HUB_REVISIONS,
|
| 18 |
+
ArtifactIntegrityError,
|
| 19 |
+
pinned_sha256,
|
| 20 |
+
resolve_hub_commit,
|
| 21 |
+
resolve_weights,
|
| 22 |
+
sha256_file,
|
| 23 |
+
)
|
| 24 |
+
|
| 25 |
+
PRODUCTION_REPO = "apoorvrajdev/captioning-inceptionv3-transformer"
|
| 26 |
+
COMMIT = "0123456789abcdef0123456789abcdef01234567"
|
| 27 |
+
UNPINNED_COMMIT = "fedcba9876543210fedcba9876543210fedcba98"
|
| 28 |
+
V1_COMMIT = "d917a47decc1f352e50ba4226e92ca5f00d47e9a"
|
| 29 |
+
V2_COMMIT = "59d93b4babb16b0ac81eef598f3abc271a355cbf"
|
| 30 |
+
|
| 31 |
+
# Distinct bytes per file, so a swapped or truncated file changes its digest.
|
| 32 |
+
ARTIFACT_BYTES = {
|
| 33 |
+
"model.h5": b"\x89HDF\r\n\x1a\nfake weights",
|
| 34 |
+
"captioning.weights.h5": b"\x89HDF\r\n\x1a\nother fake weights",
|
| 35 |
+
"vocab.json": json.dumps(["", "[UNK]", "a", "dog"]).encode(),
|
| 36 |
+
}
|
| 37 |
+
|
| 38 |
+
|
| 39 |
+
def _sha256(data: bytes) -> str:
|
| 40 |
+
return hashlib.sha256(data).hexdigest()
|
| 41 |
+
|
| 42 |
+
|
| 43 |
+
@pytest.fixture(autouse=True)
|
| 44 |
+
def _pin_test_commit(monkeypatch: pytest.MonkeyPatch) -> None:
|
| 45 |
+
"""Pin ``COMMIT`` of the test repo to the digests of ``ARTIFACT_BYTES``."""
|
| 46 |
+
monkeypatch.setitem(
|
| 47 |
+
PINNED_ARTIFACT_SHA256,
|
| 48 |
+
("user/captioning-weights", COMMIT),
|
| 49 |
+
{name: _sha256(data) for name, data in ARTIFACT_BYTES.items()},
|
| 50 |
+
)
|
| 51 |
+
|
| 52 |
+
|
| 53 |
+
def _snapshot(directory: Path, *names: str) -> Path:
|
| 54 |
+
"""Create a fake snapshot directory holding the given files' test bytes."""
|
| 55 |
+
directory.mkdir(parents=True)
|
| 56 |
+
for name in names:
|
| 57 |
+
(directory / name).write_bytes(ARTIFACT_BYTES[name])
|
| 58 |
+
return directory
|
| 59 |
+
|
| 60 |
+
|
| 61 |
+
def _flip_last_byte(path: Path) -> None:
|
| 62 |
+
data = bytearray(path.read_bytes())
|
| 63 |
+
data[-1] ^= 0x01
|
| 64 |
+
path.write_bytes(bytes(data))
|
| 65 |
+
|
| 66 |
+
|
| 67 |
+
def _recording_downloader(snapshot: Path, calls: list[dict[str, object]]):
|
| 68 |
+
def fake_downloader(
|
| 69 |
+
*, repo_id: str, revision: str, cache_dir: str | None, allow_patterns: list[str]
|
| 70 |
+
) -> str:
|
| 71 |
+
calls.append(
|
| 72 |
+
{
|
| 73 |
+
"repo_id": repo_id,
|
| 74 |
+
"revision": revision,
|
| 75 |
+
"cache_dir": cache_dir,
|
| 76 |
+
"allow_patterns": allow_patterns,
|
| 77 |
+
}
|
| 78 |
+
)
|
| 79 |
+
return str(snapshot)
|
| 80 |
+
|
| 81 |
+
return fake_downloader
|
| 82 |
+
|
| 83 |
+
|
| 84 |
+
def test_resolve_weights_local_mode_returns_settings_paths_verbatim() -> None:
|
| 85 |
+
settings = BackendSettings(
|
| 86 |
+
weights_path=Path("models/v1.0.0/model.h5"),
|
| 87 |
+
tokenizer_dir=Path("models/v1.0.0"),
|
| 88 |
+
weights_hub_repo=None,
|
| 89 |
+
)
|
| 90 |
+
|
| 91 |
+
weights_path, tokenizer_dir = resolve_weights(settings, downloader=None)
|
| 92 |
+
|
| 93 |
+
assert weights_path == Path("models/v1.0.0/model.h5")
|
| 94 |
+
assert tokenizer_dir == Path("models/v1.0.0")
|
| 95 |
+
|
| 96 |
+
|
| 97 |
+
def test_resolve_weights_hub_mode_calls_downloader_with_expected_args(tmp_path: Path) -> None:
|
| 98 |
+
fake_snapshot = _snapshot(tmp_path / "snapshots" / COMMIT, "model.h5", "vocab.json")
|
| 99 |
+
calls: list[dict[str, object]] = []
|
| 100 |
+
|
| 101 |
+
settings = BackendSettings(
|
| 102 |
+
weights_hub_repo="user/captioning-weights",
|
| 103 |
+
weights_hub_revision=COMMIT,
|
| 104 |
+
weights_hub_filename="model.h5",
|
| 105 |
+
weights_cache_dir=tmp_path / "cache",
|
| 106 |
+
)
|
| 107 |
+
|
| 108 |
+
weights_path, tokenizer_dir = resolve_weights(
|
| 109 |
+
settings, downloader=_recording_downloader(fake_snapshot, calls)
|
| 110 |
+
)
|
| 111 |
+
|
| 112 |
+
assert calls == [
|
| 113 |
+
{
|
| 114 |
+
"repo_id": "user/captioning-weights",
|
| 115 |
+
"revision": COMMIT,
|
| 116 |
+
"cache_dir": str(tmp_path / "cache"),
|
| 117 |
+
"allow_patterns": ["model.h5", "vocab.json"],
|
| 118 |
+
}
|
| 119 |
+
]
|
| 120 |
+
assert weights_path == fake_snapshot / "model.h5"
|
| 121 |
+
assert tokenizer_dir == fake_snapshot
|
| 122 |
+
|
| 123 |
+
|
| 124 |
+
def test_resolve_weights_hub_mode_passes_none_cache_dir_when_unset(tmp_path: Path) -> None:
|
| 125 |
+
fake_snapshot = _snapshot(tmp_path / "snap", "model.h5", "vocab.json")
|
| 126 |
+
calls: list[dict[str, object]] = []
|
| 127 |
+
|
| 128 |
+
settings = BackendSettings(
|
| 129 |
+
weights_hub_repo="user/captioning-weights",
|
| 130 |
+
weights_hub_revision=COMMIT,
|
| 131 |
+
weights_cache_dir=None,
|
| 132 |
+
)
|
| 133 |
+
|
| 134 |
+
resolve_weights(settings, downloader=_recording_downloader(fake_snapshot, calls))
|
| 135 |
+
|
| 136 |
+
assert [call["cache_dir"] for call in calls] == [None]
|
| 137 |
+
|
| 138 |
+
|
| 139 |
+
def test_resolve_weights_hub_mode_honors_custom_weights_filename(tmp_path: Path) -> None:
|
| 140 |
+
fake_snapshot = _snapshot(tmp_path / "snap", "captioning.weights.h5", "vocab.json")
|
| 141 |
+
calls: list[dict[str, object]] = []
|
| 142 |
+
|
| 143 |
+
settings = BackendSettings(
|
| 144 |
+
weights_hub_repo="user/captioning-weights",
|
| 145 |
+
weights_hub_revision=COMMIT,
|
| 146 |
+
weights_hub_filename="captioning.weights.h5",
|
| 147 |
+
)
|
| 148 |
+
|
| 149 |
+
weights_path, _ = resolve_weights(
|
| 150 |
+
settings, downloader=_recording_downloader(fake_snapshot, calls)
|
| 151 |
+
)
|
| 152 |
+
|
| 153 |
+
assert weights_path == fake_snapshot / "captioning.weights.h5"
|
| 154 |
+
assert calls[0]["allow_patterns"] == ["captioning.weights.h5", "vocab.json"]
|
| 155 |
+
|
| 156 |
+
|
| 157 |
+
def test_pinned_tag_downloads_its_recorded_commit_not_the_tag(tmp_path: Path) -> None:
|
| 158 |
+
"""The Space's ``v2.0.0`` variable keeps working, but the tag never reaches the Hub.
|
| 159 |
+
|
| 160 |
+
The stub serves test bytes, so the real ``v2.0.0`` digests reject them
|
| 161 |
+
after the download.
|
| 162 |
+
"""
|
| 163 |
+
fake_snapshot = _snapshot(tmp_path / "snap", "model.h5", "vocab.json")
|
| 164 |
+
calls: list[dict[str, object]] = []
|
| 165 |
+
|
| 166 |
+
settings = BackendSettings(weights_hub_repo=PRODUCTION_REPO, weights_hub_revision="v2.0.0")
|
| 167 |
+
|
| 168 |
+
with pytest.raises(ArtifactIntegrityError):
|
| 169 |
+
resolve_weights(settings, downloader=_recording_downloader(fake_snapshot, calls))
|
| 170 |
+
|
| 171 |
+
assert [call["revision"] for call in calls] == [V2_COMMIT]
|
| 172 |
+
assert [call["allow_patterns"] for call in calls] == [["model.h5", "vocab.json"]]
|
| 173 |
+
|
| 174 |
+
|
| 175 |
+
@pytest.mark.parametrize(
|
| 176 |
+
("repo_id", "revision"),
|
| 177 |
+
[
|
| 178 |
+
("user/captioning-weights", "main"), # the settings default: a branch
|
| 179 |
+
(PRODUCTION_REPO, "main"),
|
| 180 |
+
(PRODUCTION_REPO, "v3.0.0"), # a tag nobody pinned
|
| 181 |
+
(PRODUCTION_REPO, "59d93b4"), # a short commit id
|
| 182 |
+
(PRODUCTION_REPO, "59D93B4BABB16B0AC81EEF598F3ABC271A355CBF"), # not the Hub's form
|
| 183 |
+
("someone-else/captioning-weights", "v2.0.0"), # pinned only for the production repo
|
| 184 |
+
(PRODUCTION_REPO, ""),
|
| 185 |
+
],
|
| 186 |
+
)
|
| 187 |
+
def test_unpinned_revisions_are_refused_before_any_download(
|
| 188 |
+
tmp_path: Path, repo_id: str, revision: str
|
| 189 |
+
) -> None:
|
| 190 |
+
calls: list[dict[str, object]] = []
|
| 191 |
+
settings = BackendSettings(weights_hub_repo=repo_id, weights_hub_revision=revision)
|
| 192 |
+
|
| 193 |
+
with pytest.raises(ValueError, match="neither a full 40-character commit id nor a tag"):
|
| 194 |
+
resolve_weights(settings, downloader=_recording_downloader(tmp_path, calls))
|
| 195 |
+
|
| 196 |
+
assert calls == []
|
| 197 |
+
|
| 198 |
+
|
| 199 |
+
@pytest.mark.parametrize(
|
| 200 |
+
("present", "missing"),
|
| 201 |
+
[(["model.h5"], "vocab.json"), (["vocab.json"], "model.h5"), ([], "model.h5, vocab.json")],
|
| 202 |
+
)
|
| 203 |
+
def test_missing_artifacts_at_the_commit_fail_by_name(
|
| 204 |
+
tmp_path: Path, present: list[str], missing: str
|
| 205 |
+
) -> None:
|
| 206 |
+
fake_snapshot = _snapshot(tmp_path / "snap", *present)
|
| 207 |
+
settings = BackendSettings(
|
| 208 |
+
weights_hub_repo="user/captioning-weights", weights_hub_revision=COMMIT
|
| 209 |
+
)
|
| 210 |
+
|
| 211 |
+
with pytest.raises(
|
| 212 |
+
FileNotFoundError,
|
| 213 |
+
match=f"Hub commit {COMMIT} of user/captioning-weights has no {missing}\\.",
|
| 214 |
+
):
|
| 215 |
+
resolve_weights(settings, downloader=_recording_downloader(fake_snapshot, []))
|
| 216 |
+
|
| 217 |
+
|
| 218 |
+
def test_every_pinned_revision_is_itself_a_valid_commit_id() -> None:
|
| 219 |
+
assert PINNED_HUB_REVISIONS
|
| 220 |
+
for (repo_id, tag), commit in PINNED_HUB_REVISIONS.items():
|
| 221 |
+
assert resolve_hub_commit(repo_id, commit) == commit
|
| 222 |
+
assert resolve_hub_commit(repo_id, tag) == commit
|
| 223 |
+
|
| 224 |
+
|
| 225 |
+
# ---- SHA-256 verification (TASK-029) ---------------------------------------
|
| 226 |
+
|
| 227 |
+
|
| 228 |
+
def test_sha256_file_matches_hashlib_across_chunk_boundaries(tmp_path: Path) -> None:
|
| 229 |
+
data = bytes(range(256)) * 41 # 10,496 bytes: not a multiple of the chunk size
|
| 230 |
+
path = tmp_path / "blob"
|
| 231 |
+
path.write_bytes(data)
|
| 232 |
+
(tmp_path / "empty").write_bytes(b"")
|
| 233 |
+
|
| 234 |
+
assert sha256_file(path, chunk_size=1000) == hashlib.sha256(data).hexdigest()
|
| 235 |
+
assert sha256_file(path) == hashlib.sha256(data).hexdigest()
|
| 236 |
+
assert sha256_file(tmp_path / "empty") == hashlib.sha256(b"").hexdigest()
|
| 237 |
+
|
| 238 |
+
|
| 239 |
+
@pytest.mark.parametrize("tampered", ["model.h5", "vocab.json"])
|
| 240 |
+
def test_a_one_byte_change_in_either_file_is_rejected_by_name(
|
| 241 |
+
tmp_path: Path, tampered: str
|
| 242 |
+
) -> None:
|
| 243 |
+
fake_snapshot = _snapshot(tmp_path / "snap", "model.h5", "vocab.json")
|
| 244 |
+
_flip_last_byte(fake_snapshot / tampered)
|
| 245 |
+
intact = "vocab.json" if tampered == "model.h5" else "model.h5"
|
| 246 |
+
settings = BackendSettings(
|
| 247 |
+
weights_hub_repo="user/captioning-weights", weights_hub_revision=COMMIT
|
| 248 |
+
)
|
| 249 |
+
|
| 250 |
+
with pytest.raises(ArtifactIntegrityError) as excinfo:
|
| 251 |
+
resolve_weights(settings, downloader=_recording_downloader(fake_snapshot, []))
|
| 252 |
+
|
| 253 |
+
message = str(excinfo.value)
|
| 254 |
+
actual = _sha256((fake_snapshot / tampered).read_bytes())
|
| 255 |
+
expected = _sha256(ARTIFACT_BYTES[tampered])
|
| 256 |
+
assert f"{tampered} has SHA-256 {actual}, expected {expected}" in message
|
| 257 |
+
assert f"Hub commit {COMMIT} of user/captioning-weights failed verification" in message
|
| 258 |
+
assert "Nothing was loaded" in message
|
| 259 |
+
assert f"{intact} has SHA-256" not in message
|
| 260 |
+
|
| 261 |
+
|
| 262 |
+
def test_every_mismatching_file_is_reported_not_just_the_first(tmp_path: Path) -> None:
|
| 263 |
+
fake_snapshot = _snapshot(tmp_path / "snap", "model.h5", "vocab.json")
|
| 264 |
+
_flip_last_byte(fake_snapshot / "model.h5")
|
| 265 |
+
_flip_last_byte(fake_snapshot / "vocab.json")
|
| 266 |
+
settings = BackendSettings(
|
| 267 |
+
weights_hub_repo="user/captioning-weights", weights_hub_revision=COMMIT
|
| 268 |
+
)
|
| 269 |
+
|
| 270 |
+
with pytest.raises(ArtifactIntegrityError, match="model.h5 has SHA-256 .*vocab.json has"):
|
| 271 |
+
resolve_weights(settings, downloader=_recording_downloader(fake_snapshot, []))
|
| 272 |
+
|
| 273 |
+
|
| 274 |
+
def test_a_file_replaced_by_other_pinned_bytes_is_rejected(tmp_path: Path) -> None:
|
| 275 |
+
"""Each file is checked against its own digest, not against any pinned digest."""
|
| 276 |
+
fake_snapshot = _snapshot(tmp_path / "snap", "model.h5", "vocab.json")
|
| 277 |
+
(fake_snapshot / "model.h5").write_bytes(ARTIFACT_BYTES["captioning.weights.h5"])
|
| 278 |
+
settings = BackendSettings(
|
| 279 |
+
weights_hub_repo="user/captioning-weights", weights_hub_revision=COMMIT
|
| 280 |
+
)
|
| 281 |
+
|
| 282 |
+
with pytest.raises(ArtifactIntegrityError, match="model.h5 has SHA-256"):
|
| 283 |
+
resolve_weights(settings, downloader=_recording_downloader(fake_snapshot, []))
|
| 284 |
+
|
| 285 |
+
|
| 286 |
+
@pytest.mark.parametrize(
|
| 287 |
+
("repo_id", "revision"),
|
| 288 |
+
[
|
| 289 |
+
("user/captioning-weights", UNPINNED_COMMIT), # a valid commit id with no digests
|
| 290 |
+
(PRODUCTION_REPO, UNPINNED_COMMIT),
|
| 291 |
+
(PRODUCTION_REPO, COMMIT), # pinned only for the test repo
|
| 292 |
+
("someone-else/captioning-weights", V2_COMMIT), # pinned only for the production repo
|
| 293 |
+
],
|
| 294 |
+
)
|
| 295 |
+
def test_commits_without_pinned_digests_are_refused_before_any_download(
|
| 296 |
+
tmp_path: Path, repo_id: str, revision: str
|
| 297 |
+
) -> None:
|
| 298 |
+
calls: list[dict[str, object]] = []
|
| 299 |
+
settings = BackendSettings(weights_hub_repo=repo_id, weights_hub_revision=revision)
|
| 300 |
+
|
| 301 |
+
with pytest.raises(ValueError, match=f"Hub commit {revision} of {repo_id} has no SHA-256"):
|
| 302 |
+
resolve_weights(settings, downloader=_recording_downloader(tmp_path, calls))
|
| 303 |
+
|
| 304 |
+
assert calls == []
|
| 305 |
+
|
| 306 |
+
|
| 307 |
+
def test_a_pinned_tag_whose_commit_lost_its_digests_is_refused(
|
| 308 |
+
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
| 309 |
+
) -> None:
|
| 310 |
+
monkeypatch.delitem(PINNED_ARTIFACT_SHA256, (PRODUCTION_REPO, V2_COMMIT))
|
| 311 |
+
calls: list[dict[str, object]] = []
|
| 312 |
+
settings = BackendSettings(weights_hub_repo=PRODUCTION_REPO, weights_hub_revision="v2.0.0")
|
| 313 |
+
|
| 314 |
+
with pytest.raises(ValueError, match=f"Hub commit {V2_COMMIT} .* has no SHA-256 digests"):
|
| 315 |
+
resolve_weights(settings, downloader=_recording_downloader(tmp_path, calls))
|
| 316 |
+
|
| 317 |
+
assert calls == []
|
| 318 |
+
|
| 319 |
+
|
| 320 |
+
@pytest.mark.parametrize(
|
| 321 |
+
("entry", "match"),
|
| 322 |
+
[
|
| 323 |
+
({"model.h5": _sha256(b"x")}, "no SHA-256 for vocab.json"),
|
| 324 |
+
({"vocab.json": _sha256(b"x")}, "no SHA-256 for model.h5"),
|
| 325 |
+
({}, "no SHA-256 for model.h5"),
|
| 326 |
+
({"model.h5": _sha256(b"x").upper(), "vocab.json": _sha256(b"x")}, "malformed"),
|
| 327 |
+
({"model.h5": _sha256(b"x")[:-1], "vocab.json": _sha256(b"x")}, "malformed"),
|
| 328 |
+
({"model.h5": "g" * 64, "vocab.json": _sha256(b"x")}, "malformed"),
|
| 329 |
+
({"model.h5": _sha256(b"x"), "vocab.json": ""}, "malformed"),
|
| 330 |
+
],
|
| 331 |
+
)
|
| 332 |
+
def test_incomplete_or_malformed_digest_entries_are_refused_before_any_download(
|
| 333 |
+
tmp_path: Path, monkeypatch: pytest.MonkeyPatch, entry: dict[str, str], match: str
|
| 334 |
+
) -> None:
|
| 335 |
+
monkeypatch.setitem(PINNED_ARTIFACT_SHA256, ("user/captioning-weights", COMMIT), entry)
|
| 336 |
+
calls: list[dict[str, object]] = []
|
| 337 |
+
settings = BackendSettings(
|
| 338 |
+
weights_hub_repo="user/captioning-weights", weights_hub_revision=COMMIT
|
| 339 |
+
)
|
| 340 |
+
|
| 341 |
+
with pytest.raises(ValueError, match=match):
|
| 342 |
+
resolve_weights(settings, downloader=_recording_downloader(tmp_path, calls))
|
| 343 |
+
|
| 344 |
+
assert calls == []
|
| 345 |
+
|
| 346 |
+
|
| 347 |
+
def test_a_weights_filename_without_a_digest_is_refused_before_any_download(
|
| 348 |
+
tmp_path: Path,
|
| 349 |
+
) -> None:
|
| 350 |
+
calls: list[dict[str, object]] = []
|
| 351 |
+
settings = BackendSettings(
|
| 352 |
+
weights_hub_repo="user/captioning-weights",
|
| 353 |
+
weights_hub_revision=COMMIT,
|
| 354 |
+
weights_hub_filename="model.keras",
|
| 355 |
+
)
|
| 356 |
+
|
| 357 |
+
with pytest.raises(ValueError, match="no SHA-256 for model.keras"):
|
| 358 |
+
resolve_weights(settings, downloader=_recording_downloader(tmp_path, calls))
|
| 359 |
+
|
| 360 |
+
assert calls == []
|
| 361 |
+
|
| 362 |
+
|
| 363 |
+
def test_every_pinned_tag_has_well_formed_digests_for_both_served_files() -> None:
|
| 364 |
+
"""The production tag and the rollback tag both stay servable."""
|
| 365 |
+
assert {tag for _, tag in PINNED_HUB_REVISIONS} >= {"v1.0.0", "v2.0.0"}
|
| 366 |
+
digests: list[str] = []
|
| 367 |
+
for (repo_id, _tag), commit in PINNED_HUB_REVISIONS.items():
|
| 368 |
+
expected = pinned_sha256(repo_id, commit, ["model.h5", "vocab.json"])
|
| 369 |
+
assert sorted(expected) == ["model.h5", "vocab.json"]
|
| 370 |
+
digests.extend(expected.values())
|
| 371 |
+
assert len(set(digests)) == len(digests) # no revision can pass for another
|
| 372 |
+
|
| 373 |
+
|
| 374 |
+
def test_rollback_vocabulary_digest_matches_the_copy_committed_in_this_repo() -> None:
|
| 375 |
+
"""Ties one pinned digest to bytes committed here long before the pin.
|
| 376 |
+
|
| 377 |
+
Git stores this file with LF line endings. A Windows checkout with
|
| 378 |
+
``core.autocrlf`` writes CRLF, so that conversion is undone first.
|
| 379 |
+
"""
|
| 380 |
+
committed = Path(__file__).resolve().parents[3] / "models" / "v1.0.0" / "vocab.json"
|
| 381 |
+
data = committed.read_bytes().replace(b"\r\n", b"\n")
|
| 382 |
+
|
| 383 |
+
assert _sha256(data) == PINNED_ARTIFACT_SHA256[(PRODUCTION_REPO, V1_COMMIT)]["vocab.json"]
|
| 384 |
+
|
| 385 |
+
|
| 386 |
+
@pytest.mark.parametrize(("tag", "commit"), [("v1.0.0", V1_COMMIT), ("v2.0.0", V2_COMMIT)])
|
| 387 |
+
def test_each_pinned_tag_is_checked_against_its_own_commit_digests(
|
| 388 |
+
tmp_path: Path, tag: str, commit: str
|
| 389 |
+
) -> None:
|
| 390 |
+
fake_snapshot = _snapshot(tmp_path / "snap", "model.h5", "vocab.json")
|
| 391 |
+
settings = BackendSettings(weights_hub_repo=PRODUCTION_REPO, weights_hub_revision=tag)
|
| 392 |
+
|
| 393 |
+
with pytest.raises(ArtifactIntegrityError) as excinfo:
|
| 394 |
+
resolve_weights(settings, downloader=_recording_downloader(fake_snapshot, []))
|
| 395 |
+
|
| 396 |
+
message = str(excinfo.value)
|
| 397 |
+
own = PINNED_ARTIFACT_SHA256[(PRODUCTION_REPO, commit)]
|
| 398 |
+
assert f"Hub commit {commit} of {PRODUCTION_REPO}" in message
|
| 399 |
+
assert f"expected {own['model.h5']}" in message
|
| 400 |
+
assert f"expected {own['vocab.json']}" in message
|
backend/app/utils/__init__.py
ADDED
|
@@ -0,0 +1,5 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""HTTP-layer utilities (image decoding, etc.)."""
|
| 2 |
+
|
| 3 |
+
from app.utils.image import ALLOWED_CONTENT_TYPES, ImageDecodeError, bytes_to_tensor
|
| 4 |
+
|
| 5 |
+
__all__ = ["ALLOWED_CONTENT_TYPES", "ImageDecodeError", "bytes_to_tensor"]
|
backend/app/utils/image.py
ADDED
|
@@ -0,0 +1,147 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Image-decoding utilities for the HTTP boundary.
|
| 2 |
+
|
| 3 |
+
The ML package's ``inference/image_loader.py`` reads from disk; the API
|
| 4 |
+
receives bytes in memory from a multipart upload. This module bridges the
|
| 5 |
+
two: it decodes raw bytes and runs them through the *same*
|
| 6 |
+
``preprocess_image_tensor`` the training pipeline uses, so train/serve
|
| 7 |
+
parity is preserved by construction.
|
| 8 |
+
|
| 9 |
+
Before anything decodes, ``check_image_dimensions`` reads the image header
|
| 10 |
+
and rejects an image too large to decode safely: the compressed upload cap
|
| 11 |
+
doesn't bound the decoded size.
|
| 12 |
+
|
| 13 |
+
TensorFlow imports are deferred until first call to keep app import cheap
|
| 14 |
+
(e.g. when running ``ruff`` or constructing the app for tests with stub
|
| 15 |
+
predictors).
|
| 16 |
+
"""
|
| 17 |
+
|
| 18 |
+
from __future__ import annotations
|
| 19 |
+
|
| 20 |
+
import io
|
| 21 |
+
from typing import Any
|
| 22 |
+
|
| 23 |
+
from PIL import Image
|
| 24 |
+
|
| 25 |
+
# WebP isn't listed: TF 2.15's ``decode_image`` can't decode it, and serving
|
| 26 |
+
# must decode through the training path, so a WebP upload could never succeed.
|
| 27 |
+
ALLOWED_CONTENT_TYPES: frozenset[str] = frozenset(
|
| 28 |
+
{
|
| 29 |
+
"image/jpeg",
|
| 30 |
+
"image/jpg",
|
| 31 |
+
"image/png",
|
| 32 |
+
"image/bmp",
|
| 33 |
+
}
|
| 34 |
+
)
|
| 35 |
+
|
| 36 |
+
# Upload limits on decoded size, checked against the header before any decode.
|
| 37 |
+
# The model only sees 299x299, so they bound memory, not caption quality.
|
| 38 |
+
# Decoding allocates 3 bytes per pixel (uint8 RGB) and the Resizing layer then
|
| 39 |
+
# makes a float32 copy (12 bytes per pixel), so a 40 MP image needs roughly
|
| 40 |
+
# 600 MB per request. That figure is derived from the dtypes, not measured.
|
| 41 |
+
# 40 MP admits 12-36 MP camera and phone photos; the 16,384 px side admits
|
| 42 |
+
# panoramas while bounding per-row decoder buffers.
|
| 43 |
+
MAX_IMAGE_PIXELS = 40_000_000
|
| 44 |
+
MAX_IMAGE_SIDE = 16_384
|
| 45 |
+
|
| 46 |
+
# Formats whose header is read. TF 2.15's decode_image also takes GIF, but it
|
| 47 |
+
# sizes a GIF by its largest frame, which the header read doesn't see, so GIF
|
| 48 |
+
# bytes are rejected (GIF isn't on the allow-list). WebP is never parsed:
|
| 49 |
+
# Pillow's WebP open allocates the full canvas, and TF 2.15 can't decode WebP.
|
| 50 |
+
_HEADER_FORMATS = ("JPEG", "PNG", "BMP")
|
| 51 |
+
|
| 52 |
+
|
| 53 |
+
class ImageDecodeError(ValueError):
|
| 54 |
+
"""Raised when uploaded bytes are not a recognisable image."""
|
| 55 |
+
|
| 56 |
+
|
| 57 |
+
class ImageTooLargeError(ImageDecodeError):
|
| 58 |
+
"""Raised when an image's dimensions exceed the upload limits.
|
| 59 |
+
|
| 60 |
+
Subclasses ``ImageDecodeError`` so a handler that only catches that class
|
| 61 |
+
still answers with a client error, never a 500.
|
| 62 |
+
"""
|
| 63 |
+
|
| 64 |
+
|
| 65 |
+
def check_image_dimensions(
|
| 66 |
+
image_bytes: bytes,
|
| 67 |
+
*,
|
| 68 |
+
max_side: int = MAX_IMAGE_SIDE,
|
| 69 |
+
max_pixels: int = MAX_IMAGE_PIXELS,
|
| 70 |
+
) -> tuple[int, int]:
|
| 71 |
+
"""Read the image header and reject an image too large to decode safely.
|
| 72 |
+
|
| 73 |
+
``tf.io.decode_image`` allocates the full-resolution tensor from the
|
| 74 |
+
header's dimensions before it reads any pixel data, so a small file can
|
| 75 |
+
claim a huge image. ``Image.open`` parses only the header here; no pixel
|
| 76 |
+
is decoded.
|
| 77 |
+
|
| 78 |
+
Args:
|
| 79 |
+
image_bytes: Raw upload bytes.
|
| 80 |
+
max_side: Largest width or height accepted, in pixels.
|
| 81 |
+
max_pixels: Largest ``width * height`` accepted.
|
| 82 |
+
|
| 83 |
+
Returns:
|
| 84 |
+
``(width, height)`` read from the header.
|
| 85 |
+
|
| 86 |
+
Raises:
|
| 87 |
+
ImageTooLargeError: If either side or the pixel count is over its limit.
|
| 88 |
+
ImageDecodeError: If the bytes aren't a JPEG, PNG or BMP with a
|
| 89 |
+
readable header and positive dimensions.
|
| 90 |
+
"""
|
| 91 |
+
too_large = (
|
| 92 |
+
f"Image dimensions exceed the limit of {max_side} pixels per side "
|
| 93 |
+
f"and {max_pixels:,} pixels in total."
|
| 94 |
+
)
|
| 95 |
+
try:
|
| 96 |
+
with Image.open(io.BytesIO(image_bytes), formats=_HEADER_FORMATS) as image:
|
| 97 |
+
width, height = image.size
|
| 98 |
+
except Image.DecompressionBombError as exc:
|
| 99 |
+
# Pillow's own ceiling (about 179 MP) is far above max_pixels.
|
| 100 |
+
raise ImageTooLargeError(too_large) from exc
|
| 101 |
+
except Exception as exc:
|
| 102 |
+
# Pillow's header parsers raise many exception types on hostile input.
|
| 103 |
+
# Each one means the upload isn't an image this service accepts.
|
| 104 |
+
raise ImageDecodeError(
|
| 105 |
+
"Could not decode image bytes: expected a JPEG, PNG or BMP image."
|
| 106 |
+
) from exc
|
| 107 |
+
|
| 108 |
+
if width <= 0 or height <= 0:
|
| 109 |
+
raise ImageDecodeError("Could not decode image bytes: invalid image dimensions.")
|
| 110 |
+
if width > max_side or height > max_side or width * height > max_pixels:
|
| 111 |
+
raise ImageTooLargeError(too_large)
|
| 112 |
+
return width, height
|
| 113 |
+
|
| 114 |
+
|
| 115 |
+
def bytes_to_tensor(image_bytes: bytes) -> Any:
|
| 116 |
+
"""Decode an in-memory image into a model-ready tensor.
|
| 117 |
+
|
| 118 |
+
Args:
|
| 119 |
+
image_bytes: Raw bytes from a multipart upload (JPEG/PNG/BMP).
|
| 120 |
+
|
| 121 |
+
Returns:
|
| 122 |
+
``tf.Tensor`` of shape ``[299, 299, 3]``, dtype ``float32``, with
|
| 123 |
+
the InceptionV3 normalisation applied — i.e. exactly what
|
| 124 |
+
``CaptionPredictor.predict_tensor`` expects.
|
| 125 |
+
|
| 126 |
+
Raises:
|
| 127 |
+
ImageTooLargeError: If the header dimensions exceed the limits.
|
| 128 |
+
ImageDecodeError: If the bytes can't be decoded as an image.
|
| 129 |
+
"""
|
| 130 |
+
# Runs before TensorFlow is even imported: the decode below allocates the
|
| 131 |
+
# whole image as soon as it has read the header.
|
| 132 |
+
check_image_dimensions(image_bytes)
|
| 133 |
+
|
| 134 |
+
import tensorflow as tf
|
| 135 |
+
|
| 136 |
+
from captioning.preprocessing.image import preprocess_image_tensor
|
| 137 |
+
|
| 138 |
+
try:
|
| 139 |
+
decoded = tf.io.decode_image(
|
| 140 |
+
image_bytes,
|
| 141 |
+
channels=3,
|
| 142 |
+
expand_animations=False,
|
| 143 |
+
)
|
| 144 |
+
except (tf.errors.InvalidArgumentError, tf.errors.UnknownError) as exc:
|
| 145 |
+
raise ImageDecodeError(f"Could not decode image bytes: {exc}") from exc
|
| 146 |
+
|
| 147 |
+
return preprocess_image_tensor(decoded)
|
configs/base.yaml
ADDED
|
@@ -0,0 +1,78 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# =============================================================================
|
| 2 |
+
# configs/base.yaml — single canonical config for training and inference.
|
| 3 |
+
# -----------------------------------------------------------------------------
|
| 4 |
+
# Every value here mirrors the IEEE notebook (cell 6 hyperparams + cell 21
|
| 5 |
+
# layer wiring) so behaviour is identical to the published research. Override
|
| 6 |
+
# any field on the CLI or via env var (CAPTIONING__TRAIN__BATCH_SIZE=32) — see
|
| 7 |
+
# src/captioning/config/schema.py for the full validated schema.
|
| 8 |
+
# =============================================================================
|
| 9 |
+
|
| 10 |
+
data:
|
| 11 |
+
# Local path; scripts/prepare_data.py downloads COCO into this directory.
|
| 12 |
+
base_path: data/coco2017
|
| 13 |
+
annotations_filename: captions_train2017.json
|
| 14 |
+
images_subdir: train2017
|
| 15 |
+
sample_size: 120000 # Notebook: captions.sample(120000)
|
| 16 |
+
train_val_split: 0.8 # Notebook cell 11: int(len(img_keys) * 0.8)
|
| 17 |
+
|
| 18 |
+
model:
|
| 19 |
+
embedding_dim: 512 # Notebook: EMBEDDING_DIM = 512
|
| 20 |
+
units: 512 # Notebook: UNITS = 512
|
| 21 |
+
max_length: 40 # Notebook: MAX_LENGTH = 40
|
| 22 |
+
vocabulary_size: 15000 # Notebook: VOCABULARY_SIZE = 15000
|
| 23 |
+
encoder_num_heads: 1 # Notebook cell 21: TransformerEncoderLayer(EMBEDDING_DIM, 1)
|
| 24 |
+
decoder_num_heads: 8 # Notebook cell 21: TransformerDecoderLayer(..., 8)
|
| 25 |
+
decoder_dropout_inner: 0.3 # Notebook cell 19: dropout_1 = Dropout(0.3)
|
| 26 |
+
decoder_dropout_outer: 0.5 # Notebook cell 19: dropout_2 = Dropout(0.5)
|
| 27 |
+
decoder_attention_dropout: 0.1 # Notebook cell 19: MultiHeadAttention(dropout=0.1)
|
| 28 |
+
|
| 29 |
+
train:
|
| 30 |
+
epochs: 10 # Notebook: EPOCHS = 10
|
| 31 |
+
batch_size: 64 # Notebook: BATCH_SIZE = 64
|
| 32 |
+
buffer_size: 1000 # Notebook: BUFFER_SIZE = 1000
|
| 33 |
+
early_stopping_patience: 3 # Notebook cell 22: EarlyStopping(patience=3, ...)
|
| 34 |
+
seed: 42 # NEW: pin RNGs (notebook didn't seed; results varied)
|
| 35 |
+
learning_rate: 0.001 # Keras Adam default — what the notebook uses implicitly
|
| 36 |
+
weights_filename: model.h5 # Notebook cell 30: caption_model.save_weights('model.h5')
|
| 37 |
+
|
| 38 |
+
serve:
|
| 39 |
+
max_upload_bytes: 10485760 # 10 MB — guard at the API edge
|
| 40 |
+
decode_strategy: greedy # Phase 1b: "beam"
|
| 41 |
+
# Beam settings for requests that ask for `decode_strategy=beam` (TASK-034).
|
| 42 |
+
# They are the committed evaluation's: results/stabilized-beam-w4-lp07-rp12
|
| 43 |
+
# (run with configs/train/stabilized.yaml, whose no_repeat_ngram_size is 3).
|
| 44 |
+
beam_width: 4
|
| 45 |
+
length_penalty: 0.7
|
| 46 |
+
repetition_penalty: 1.2
|
| 47 |
+
no_repeat_ngram_size: 3
|
| 48 |
+
cors_allowed_origins:
|
| 49 |
+
- http://localhost:3000
|
| 50 |
+
- http://localhost:5173
|
| 51 |
+
- http://localhost:5174
|
| 52 |
+
- http://127.0.0.1:5173
|
| 53 |
+
- http://127.0.0.1:5174
|
| 54 |
+
|
| 55 |
+
# Phase 3 comparison — frozen protocol, docs/EVAL_METHODOLOGY.md § 8 (ADR-019).
|
| 56 |
+
# Changing a value here is a dated protocol amendment, not a tuning knob.
|
| 57 |
+
compare:
|
| 58 |
+
cnn: # § 8.1: the project's model as compared
|
| 59 |
+
model_id: inceptionv3-transformer-stabilized
|
| 60 |
+
hub_repo: apoorvrajdev/captioning-inceptionv3-transformer
|
| 61 |
+
revision: 59d93b4babb16b0ac81eef598f3abc271a355cbf # tag v2.0.0
|
| 62 |
+
baselines: # § 8.1: Hugging Face baselines, pinned commits
|
| 63 |
+
- model_id: blip-base
|
| 64 |
+
hub_repo: Salesforce/blip-image-captioning-base
|
| 65 |
+
revision: 82a37760796d32b1411fe092ab5d4e227313294b
|
| 66 |
+
- model_id: vit-gpt2
|
| 67 |
+
hub_repo: nlpconnect/vit-gpt2-image-captioning
|
| 68 |
+
revision: dc68f91c06a1ba6f15268e5b9c13ae7a7c514084
|
| 69 |
+
- model_id: git-base-coco
|
| 70 |
+
hub_repo: microsoft/git-base-coco
|
| 71 |
+
revision: a13141da42abd4a8cbf283601a8104265f537cee
|
| 72 |
+
baseline_decode: # § 8.4: passed explicitly to generate()
|
| 73 |
+
num_beams: 1
|
| 74 |
+
do_sample: false
|
| 75 |
+
max_new_tokens: 40
|
| 76 |
+
repetition_penalty: 1.0
|
| 77 |
+
no_repeat_ngram_size: 0
|
| 78 |
+
precision: float32
|
configs/train/debug.yaml
ADDED
|
@@ -0,0 +1,18 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# =============================================================================
|
| 2 |
+
# configs/train/debug.yaml — fast end-to-end smoke run.
|
| 3 |
+
# -----------------------------------------------------------------------------
|
| 4 |
+
# Used by CI to verify the training pipeline imports and steps once without
|
| 5 |
+
# OOMing or producing NaNs. Loads on top of base.yaml so only the changed
|
| 6 |
+
# fields need to be listed.
|
| 7 |
+
#
|
| 8 |
+
# python -m scripts.train --config configs/base.yaml --override configs/train/debug.yaml
|
| 9 |
+
# =============================================================================
|
| 10 |
+
|
| 11 |
+
data:
|
| 12 |
+
sample_size: 64 # Just enough captions to fill one batch
|
| 13 |
+
|
| 14 |
+
train:
|
| 15 |
+
epochs: 1
|
| 16 |
+
batch_size: 8
|
| 17 |
+
buffer_size: 16
|
| 18 |
+
seed: 0
|
configs/train/stabilized.yaml
ADDED
|
@@ -0,0 +1,84 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# =============================================================================
|
| 2 |
+
# configs/train/stabilized.yaml — first experimental run with the opt-in
|
| 3 |
+
# training-stability primitives turned on.
|
| 4 |
+
# -----------------------------------------------------------------------------
|
| 5 |
+
# Identical to configs/base.yaml except for the four flags called out in the
|
| 6 |
+
# `train:` section below. Every other field mirrors the IEEE notebook verbatim
|
| 7 |
+
# so this run is comparable to the baseline at the same seed and architecture.
|
| 8 |
+
#
|
| 9 |
+
# Why a complete config (not a thin override)?
|
| 10 |
+
# scripts/train.py only accepts --config; there is no --override merge mode
|
| 11 |
+
# in the CLI (the README mentions one but it's aspirational, not implemented).
|
| 12 |
+
# Duplicating the values here is the smallest correct change that keeps
|
| 13 |
+
# base.yaml itself untouched — which was the explicit requirement for this
|
| 14 |
+
# experiment phase.
|
| 15 |
+
#
|
| 16 |
+
# Usage:
|
| 17 |
+
# python -m scripts.train --config configs/train/stabilized.yaml \
|
| 18 |
+
# --output-dir outputs/runs/stabilized
|
| 19 |
+
#
|
| 20 |
+
# Compare against the baseline by training the same code twice — once with
|
| 21 |
+
# configs/base.yaml, once with this file — and diffing the resulting
|
| 22 |
+
# results/<run_id>/metrics.json files.
|
| 23 |
+
# =============================================================================
|
| 24 |
+
|
| 25 |
+
data:
|
| 26 |
+
base_path: data/coco2017
|
| 27 |
+
annotations_filename: captions_train2017.json
|
| 28 |
+
images_subdir: train2017
|
| 29 |
+
sample_size: 120000 # Same sample as base.yaml — comparability matters
|
| 30 |
+
train_val_split: 0.8
|
| 31 |
+
|
| 32 |
+
model:
|
| 33 |
+
embedding_dim: 512
|
| 34 |
+
units: 512
|
| 35 |
+
max_length: 40
|
| 36 |
+
vocabulary_size: 15000
|
| 37 |
+
encoder_num_heads: 1
|
| 38 |
+
decoder_num_heads: 8
|
| 39 |
+
decoder_dropout_inner: 0.3
|
| 40 |
+
decoder_dropout_outer: 0.5
|
| 41 |
+
decoder_attention_dropout: 0.1
|
| 42 |
+
|
| 43 |
+
train:
|
| 44 |
+
epochs: 10
|
| 45 |
+
batch_size: 64
|
| 46 |
+
buffer_size: 1000
|
| 47 |
+
early_stopping_patience: 3
|
| 48 |
+
seed: 42
|
| 49 |
+
learning_rate: 0.001
|
| 50 |
+
weights_filename: model.h5
|
| 51 |
+
|
| 52 |
+
# ---- the four flags this experiment is actually testing -------------------
|
| 53 |
+
# Label smoothing 0.1 softens the cross-entropy target so the decoder
|
| 54 |
+
# cannot collapse onto a handful of high-frequency tokens. Standard
|
| 55 |
+
# transformer captioning recipe (BLIP, ViT-GPT2, GIT all use it).
|
| 56 |
+
label_smoothing: 0.1
|
| 57 |
+
|
| 58 |
+
# Warmup + cosine decay replaces the bare constant Adam LR. Transformers
|
| 59 |
+
# trained from scratch with no warmup tend to settle into a "safe captions"
|
| 60 |
+
# basin where every output looks like "a man standing ...". Cosine decay
|
| 61 |
+
# then anneals smoothly toward min_learning_rate.
|
| 62 |
+
lr_schedule: cosine
|
| 63 |
+
warmup_steps: 500 # ~1/3 of an epoch at batch 64, sample 120k
|
| 64 |
+
cosine_decay_steps: null # null -> trainer derives from steps_per_epoch * epochs
|
| 65 |
+
min_learning_rate: 0.0
|
| 66 |
+
|
| 67 |
+
# Restore conventional behaviour: dropout OFF during validation, accuracy
|
| 68 |
+
# tracker weighted by token count. This gives a clean val_loss signal so
|
| 69 |
+
# EarlyStopping fires on a real plateau rather than on dropout noise.
|
| 70 |
+
honour_training_flag_in_test_step: true
|
| 71 |
+
|
| 72 |
+
serve:
|
| 73 |
+
max_upload_bytes: 10485760
|
| 74 |
+
decode_strategy: greedy # Decode strategy is selected at evaluate time
|
| 75 |
+
beam_width: 4 # Stored defaults for `scripts.evaluate --decode-strategy beam`
|
| 76 |
+
length_penalty: 0.7
|
| 77 |
+
repetition_penalty: 1.0
|
| 78 |
+
no_repeat_ngram_size: 3
|
| 79 |
+
cors_allowed_origins:
|
| 80 |
+
- http://localhost:3000
|
| 81 |
+
- http://localhost:5173
|
| 82 |
+
- http://localhost:5174
|
| 83 |
+
- http://127.0.0.1:5173
|
| 84 |
+
- http://127.0.0.1:5174
|
models/v1.0.0/vocab.json
ADDED
|
@@ -0,0 +1,54 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
[
|
| 2 |
+
"",
|
| 3 |
+
"[UNK]",
|
| 4 |
+
"a",
|
| 5 |
+
"[start]",
|
| 6 |
+
"[end]",
|
| 7 |
+
"on",
|
| 8 |
+
"of",
|
| 9 |
+
"in",
|
| 10 |
+
"wooden",
|
| 11 |
+
"table",
|
| 12 |
+
"standing",
|
| 13 |
+
"sitting",
|
| 14 |
+
"woman",
|
| 15 |
+
"with",
|
| 16 |
+
"wave",
|
| 17 |
+
"two",
|
| 18 |
+
"tree",
|
| 19 |
+
"top",
|
| 20 |
+
"surfboard",
|
| 21 |
+
"street",
|
| 22 |
+
"stove",
|
| 23 |
+
"soccer",
|
| 24 |
+
"small",
|
| 25 |
+
"riding",
|
| 26 |
+
"refrigerator",
|
| 27 |
+
"red",
|
| 28 |
+
"plate",
|
| 29 |
+
"person",
|
| 30 |
+
"people",
|
| 31 |
+
"park",
|
| 32 |
+
"mountain",
|
| 33 |
+
"man",
|
| 34 |
+
"kitchen",
|
| 35 |
+
"kicking",
|
| 36 |
+
"holding",
|
| 37 |
+
"her",
|
| 38 |
+
"group",
|
| 39 |
+
"front",
|
| 40 |
+
"food",
|
| 41 |
+
"driving",
|
| 42 |
+
"down",
|
| 43 |
+
"dog",
|
| 44 |
+
"city",
|
| 45 |
+
"child",
|
| 46 |
+
"cat",
|
| 47 |
+
"bus",
|
| 48 |
+
"branch",
|
| 49 |
+
"birds",
|
| 50 |
+
"beach",
|
| 51 |
+
"ball",
|
| 52 |
+
"arms",
|
| 53 |
+
"and"
|
| 54 |
+
]
|
pyproject.toml
ADDED
|
@@ -0,0 +1,248 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# =============================================================================
|
| 2 |
+
# pyproject.toml — single source of truth for the `captioning` Python package
|
| 3 |
+
# -----------------------------------------------------------------------------
|
| 4 |
+
# This file follows PEP 621 (project metadata) and PEP 517/518 (build system).
|
| 5 |
+
# It replaces a scattered mix of setup.py + requirements.txt + setup.cfg with
|
| 6 |
+
# one canonical config. `pip install -e .` installs the package from `src/`.
|
| 7 |
+
#
|
| 8 |
+
# Why src/ layout? It prevents accidental imports of the package from the
|
| 9 |
+
# repo root during testing — every test exercises the *installed* package,
|
| 10 |
+
# the way users will actually import it. This is the layout used by the
|
| 11 |
+
# Python Packaging Authority's example projects and recommended by pytest.
|
| 12 |
+
# =============================================================================
|
| 13 |
+
|
| 14 |
+
[build-system]
|
| 15 |
+
requires = ["setuptools>=68", "wheel"]
|
| 16 |
+
build-backend = "setuptools.build_meta"
|
| 17 |
+
|
| 18 |
+
# -----------------------------------------------------------------------------
|
| 19 |
+
# Project metadata — what `pip show captioning` will display.
|
| 20 |
+
# -----------------------------------------------------------------------------
|
| 21 |
+
[project]
|
| 22 |
+
name = "captioning"
|
| 23 |
+
version = "0.1.0"
|
| 24 |
+
description = "IEEE-published CNN+Transformer image captioning, restructured into a production-grade multimodal AI platform."
|
| 25 |
+
readme = "README.md"
|
| 26 |
+
requires-python = ">=3.10,<3.13"
|
| 27 |
+
license = { text = "MIT" }
|
| 28 |
+
authors = [
|
| 29 |
+
{ name = "Apoorv Raj" },
|
| 30 |
+
]
|
| 31 |
+
keywords = [
|
| 32 |
+
"image-captioning",
|
| 33 |
+
"multimodal",
|
| 34 |
+
"transformer",
|
| 35 |
+
"computer-vision",
|
| 36 |
+
"tensorflow",
|
| 37 |
+
"fastapi",
|
| 38 |
+
]
|
| 39 |
+
classifiers = [
|
| 40 |
+
"Development Status :: 3 - Alpha",
|
| 41 |
+
"Intended Audience :: Developers",
|
| 42 |
+
"Intended Audience :: Science/Research",
|
| 43 |
+
"License :: OSI Approved :: MIT License",
|
| 44 |
+
"Programming Language :: Python :: 3",
|
| 45 |
+
"Programming Language :: Python :: 3.10",
|
| 46 |
+
"Programming Language :: Python :: 3.11",
|
| 47 |
+
"Programming Language :: Python :: 3.12",
|
| 48 |
+
"Topic :: Scientific/Engineering :: Artificial Intelligence",
|
| 49 |
+
]
|
| 50 |
+
|
| 51 |
+
# -----------------------------------------------------------------------------
|
| 52 |
+
# Runtime dependencies for the core ML library + FastAPI backend.
|
| 53 |
+
# Pinned hard. Floating versions across TF + transformers + torch is the most
|
| 54 |
+
# common source of silent BLEU drift between training runs and deployments.
|
| 55 |
+
#
|
| 56 |
+
# - tensorflow-cpu (NOT tensorflow): saves ~600 MB and removes the CUDA driver
|
| 57 |
+
# dependency. We are deploying to CPU-only HuggingFace Spaces. If you ever
|
| 58 |
+
# train on a GPU box, install `tensorflow==2.15.0` in that environment only.
|
| 59 |
+
# - 2.15.0 specifically: TF 2.16 swapped to Keras 3 by default and broke the
|
| 60 |
+
# `tf.keras.layers.TextVectorization` saving behaviour the IEEE notebook
|
| 61 |
+
# relies on. Stay on 2.15 for v1; upgrade is a deliberate Phase-5+ task.
|
| 62 |
+
# - pydantic 2.x: required by FastAPI >= 0.100. Faster and stricter than v1.
|
| 63 |
+
# -----------------------------------------------------------------------------
|
| 64 |
+
dependencies = [
|
| 65 |
+
"tensorflow-cpu==2.15.0",
|
| 66 |
+
"numpy>=1.26,<2.0", # NumPy 2.0 broke TF 2.15 binary compat
|
| 67 |
+
"pandas>=2.1,<3.0",
|
| 68 |
+
"pillow>=12.3,<13.0", # upload header checks on the serving path; [hf] image loading
|
| 69 |
+
"pyyaml>=6.0,<7.0",
|
| 70 |
+
"pydantic>=2.7,<3.0",
|
| 71 |
+
"pydantic-settings>=2.3,<3.0",
|
| 72 |
+
"fastapi>=0.133,<1.0",
|
| 73 |
+
"starlette>=1.3.1,<2.0", # FastAPI's own floor (0.40) admits vulnerable releases
|
| 74 |
+
"uvicorn[standard]>=0.30,<1.0",
|
| 75 |
+
"python-multipart>=0.0.31", # FastAPI multipart form data (image upload)
|
| 76 |
+
"huggingface-hub>=0.23,<1.0", # Pulls weights from HF Hub at startup
|
| 77 |
+
"structlog>=24.1,<25.0", # Structured JSON logs in prod, pretty in dev
|
| 78 |
+
"anyio>=4.14.2,<5.0", # Thread-pool offload for sync TF inference
|
| 79 |
+
"tqdm>=4.66,<5.0",
|
| 80 |
+
"click>=8.1,<9.0", # CLI for scripts/
|
| 81 |
+
]
|
| 82 |
+
|
| 83 |
+
# -----------------------------------------------------------------------------
|
| 84 |
+
# Optional dependency groups — installed via `pip install ".[dev,eval]"`.
|
| 85 |
+
# Splitting these keeps the production Docker image small (Phase 1 backend
|
| 86 |
+
# image is ~1.1 GB; adding `hf` extras takes it to ~2.3 GB which is the
|
| 87 |
+
# Phase 3 comparison image).
|
| 88 |
+
# -----------------------------------------------------------------------------
|
| 89 |
+
[project.optional-dependencies]
|
| 90 |
+
|
| 91 |
+
# Tier-1 multimodal upgrade: BLIP, ViT-GPT2, GIT models from HuggingFace.
|
| 92 |
+
# torch CPU is large (~700 MB); only install when serving the comparison demo.
|
| 93 |
+
hf = [
|
| 94 |
+
"transformers==4.41.2",
|
| 95 |
+
"torch==2.3.0",
|
| 96 |
+
"sentencepiece>=0.2.0",
|
| 97 |
+
"accelerate>=0.30,<1.0",
|
| 98 |
+
]
|
| 99 |
+
|
| 100 |
+
# Evaluation metrics. Pulled separately because pycocoevalcap drags Java
|
| 101 |
+
# dependencies (METEOR), which we don't want in the serving image.
|
| 102 |
+
eval = [
|
| 103 |
+
"sacrebleu>=2.4,<3.0",
|
| 104 |
+
"nltk>=3.8,<4.0",
|
| 105 |
+
"rouge-score>=0.1.2",
|
| 106 |
+
"pycocoevalcap>=1.2",
|
| 107 |
+
"matplotlib>=3.8,<4.0",
|
| 108 |
+
]
|
| 109 |
+
|
| 110 |
+
# Experiment tracking. Local SQLite by default; points at DagsHub in prod.
|
| 111 |
+
mlflow = [
|
| 112 |
+
"mlflow>=2.13,<3.0",
|
| 113 |
+
]
|
| 114 |
+
|
| 115 |
+
# Developer tooling: lint, type-check, test. Never deployed.
|
| 116 |
+
dev = [
|
| 117 |
+
"ruff>=0.5,<1.0",
|
| 118 |
+
"mypy>=1.10,<2.0",
|
| 119 |
+
"pytest>=8.2,<9.0",
|
| 120 |
+
"pytest-cov>=5.0,<6.0",
|
| 121 |
+
"pytest-asyncio>=0.23,<1.0",
|
| 122 |
+
"httpx>=0.27,<1.0", # FastAPI TestClient backend
|
| 123 |
+
"pre-commit>=3.7,<4.0",
|
| 124 |
+
"nbstripout>=0.7,<1.0",
|
| 125 |
+
"types-PyYAML",
|
| 126 |
+
"types-requests",
|
| 127 |
+
"pandas-stubs>=2.2,<3.0",
|
| 128 |
+
]
|
| 129 |
+
|
| 130 |
+
# -----------------------------------------------------------------------------
|
| 131 |
+
# Where pip should install the package from (the `src/` layout).
|
| 132 |
+
# -----------------------------------------------------------------------------
|
| 133 |
+
[tool.setuptools.packages.find]
|
| 134 |
+
where = ["src"]
|
| 135 |
+
include = ["captioning*"]
|
| 136 |
+
|
| 137 |
+
[tool.setuptools.package-data]
|
| 138 |
+
"captioning" = ["py.typed"] # PEP 561: ship type hints with the package
|
| 139 |
+
|
| 140 |
+
# =============================================================================
|
| 141 |
+
# Tooling configuration — co-located so a single file owns project policy.
|
| 142 |
+
# =============================================================================
|
| 143 |
+
|
| 144 |
+
# ---- Ruff: linter + formatter (replaces black + isort + flake8) -------------
|
| 145 |
+
# We prefer Ruff because it runs ~100x faster and is the de-facto modern
|
| 146 |
+
# default in the Python ecosystem. One tool, one config, one cache.
|
| 147 |
+
[tool.ruff]
|
| 148 |
+
line-length = 100
|
| 149 |
+
target-version = "py310"
|
| 150 |
+
src = ["src", "backend", "scripts", "tests"]
|
| 151 |
+
extend-exclude = [
|
| 152 |
+
"notebooks", # Notebooks have their own conventions
|
| 153 |
+
"outputs",
|
| 154 |
+
"mlruns",
|
| 155 |
+
"frontend",
|
| 156 |
+
]
|
| 157 |
+
|
| 158 |
+
[tool.ruff.lint]
|
| 159 |
+
# Curated rule set — pragmatic defaults, not the full strict catalogue.
|
| 160 |
+
select = [
|
| 161 |
+
"E", # pycodestyle errors
|
| 162 |
+
"W", # pycodestyle warnings
|
| 163 |
+
"F", # pyflakes
|
| 164 |
+
"I", # isort import sorting
|
| 165 |
+
"B", # flake8-bugbear (likely bugs)
|
| 166 |
+
"UP", # pyupgrade (modern syntax)
|
| 167 |
+
"SIM", # flake8-simplify
|
| 168 |
+
"RET", # flake8-return
|
| 169 |
+
"PTH", # flake8-use-pathlib (prefer pathlib over os.path)
|
| 170 |
+
"RUF", # Ruff's own rules
|
| 171 |
+
]
|
| 172 |
+
ignore = [
|
| 173 |
+
"E501", # line length — formatter handles it; lint warnings are noise
|
| 174 |
+
"B008", # function call in default arg (FastAPI's Depends() pattern)
|
| 175 |
+
]
|
| 176 |
+
|
| 177 |
+
[tool.ruff.lint.per-file-ignores]
|
| 178 |
+
"tests/**" = ["B011"] # asserts in tests are fine
|
| 179 |
+
"scripts/**" = ["T201"] # print() in CLI scripts is fine
|
| 180 |
+
|
| 181 |
+
[tool.ruff.format]
|
| 182 |
+
quote-style = "double"
|
| 183 |
+
indent-style = "space"
|
| 184 |
+
docstring-code-format = true
|
| 185 |
+
|
| 186 |
+
# ---- Mypy: static type checker -----------------------------------------------
|
| 187 |
+
# We only enforce types on our own code; third-party untyped libs are tolerated.
|
| 188 |
+
[tool.mypy]
|
| 189 |
+
python_version = "3.10"
|
| 190 |
+
strict = false # Start lenient; tighten as types stabilise
|
| 191 |
+
warn_unused_configs = true
|
| 192 |
+
warn_redundant_casts = true
|
| 193 |
+
warn_unused_ignores = true
|
| 194 |
+
warn_no_return = true
|
| 195 |
+
no_implicit_optional = true
|
| 196 |
+
files = ["src/captioning", "backend/app", "scripts"]
|
| 197 |
+
mypy_path = ["src", "backend"]
|
| 198 |
+
explicit_package_bases = true
|
| 199 |
+
namespace_packages = true
|
| 200 |
+
|
| 201 |
+
[[tool.mypy.overrides]]
|
| 202 |
+
module = [
|
| 203 |
+
"tensorflow.*",
|
| 204 |
+
"transformers.*",
|
| 205 |
+
"huggingface_hub.*",
|
| 206 |
+
"PIL.*",
|
| 207 |
+
"nltk.*",
|
| 208 |
+
"sacrebleu.*",
|
| 209 |
+
"rouge_score.*",
|
| 210 |
+
"pycocoevalcap.*",
|
| 211 |
+
]
|
| 212 |
+
ignore_missing_imports = true
|
| 213 |
+
|
| 214 |
+
# ---- Pytest -------------------------------------------------------------------
|
| 215 |
+
[tool.pytest.ini_options]
|
| 216 |
+
minversion = "8.0"
|
| 217 |
+
testpaths = ["tests", "backend/app/tests"]
|
| 218 |
+
pythonpath = ["backend"] # Lets `from app.* import ...` resolve in tests
|
| 219 |
+
addopts = [
|
| 220 |
+
"-ra", # Show short summary for non-passing tests
|
| 221 |
+
"--strict-markers",
|
| 222 |
+
"--strict-config",
|
| 223 |
+
"--showlocals",
|
| 224 |
+
]
|
| 225 |
+
markers = [
|
| 226 |
+
"slow: tests that take >10 seconds (run with -m slow)",
|
| 227 |
+
"gpu: tests requiring a GPU (skipped in CI by default)",
|
| 228 |
+
]
|
| 229 |
+
filterwarnings = [
|
| 230 |
+
"ignore::DeprecationWarning:tensorflow.*",
|
| 231 |
+
"ignore::FutureWarning:tensorflow.*",
|
| 232 |
+
]
|
| 233 |
+
|
| 234 |
+
# ---- Coverage -----------------------------------------------------------------
|
| 235 |
+
[tool.coverage.run]
|
| 236 |
+
branch = true
|
| 237 |
+
source = ["src/captioning", "backend/app"]
|
| 238 |
+
omit = ["*/tests/*", "*/__init__.py"]
|
| 239 |
+
|
| 240 |
+
[tool.coverage.report]
|
| 241 |
+
exclude_lines = [
|
| 242 |
+
"pragma: no cover",
|
| 243 |
+
"raise NotImplementedError",
|
| 244 |
+
"if TYPE_CHECKING:",
|
| 245 |
+
"if __name__ == .__main__.:",
|
| 246 |
+
]
|
| 247 |
+
show_missing = true
|
| 248 |
+
skip_covered = false
|
requirements.txt
ADDED
|
@@ -0,0 +1,51 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# =============================================================================
|
| 2 |
+
# requirements.txt — runtime dependencies for the FastAPI backend (slim).
|
| 3 |
+
# -----------------------------------------------------------------------------
|
| 4 |
+
# This file mirrors `[project.dependencies]` in pyproject.toml so that
|
| 5 |
+
# Docker builds and CI can `pip install -r requirements.txt` without needing
|
| 6 |
+
# the package source available.
|
| 7 |
+
#
|
| 8 |
+
# To regenerate from pyproject.toml later (recommended after Phase 1):
|
| 9 |
+
# pip install pip-tools
|
| 10 |
+
# pip-compile --extra=hf pyproject.toml -o requirements.txt
|
| 11 |
+
#
|
| 12 |
+
# All versions are pinned. Floating versions across TF + transformers + torch
|
| 13 |
+
# is the most common silent source of BLEU drift between runs and deploys.
|
| 14 |
+
# =============================================================================
|
| 15 |
+
|
| 16 |
+
# ---- Core ML framework -------------------------------------------------------
|
| 17 |
+
# CPU-only TF: deploy target is CPU HuggingFace Spaces. Pinned at 2.15 because
|
| 18 |
+
# 2.16+ ships Keras 3 by default, which breaks the IEEE notebook's
|
| 19 |
+
# `tf.keras.layers.TextVectorization` save/load semantics.
|
| 20 |
+
tensorflow-cpu==2.15.0
|
| 21 |
+
|
| 22 |
+
# ---- Numerics / data ---------------------------------------------------------
|
| 23 |
+
# NumPy <2.0 because TF 2.15 was built against the 1.x ABI.
|
| 24 |
+
numpy==1.26.4
|
| 25 |
+
pandas==2.2.2
|
| 26 |
+
# Serving reads upload headers with it (JPEG, PNG and BMP only) to enforce the
|
| 27 |
+
# pixel limits; pixels still decode via tf.io.decode_image. The [hf] baseline
|
| 28 |
+
# adapter reads images with it.
|
| 29 |
+
pillow==12.3.0
|
| 30 |
+
pyyaml==6.0.1
|
| 31 |
+
|
| 32 |
+
# ---- Config & validation -----------------------------------------------------
|
| 33 |
+
pydantic==2.7.4
|
| 34 |
+
pydantic-settings==2.3.4
|
| 35 |
+
|
| 36 |
+
# ---- API server --------------------------------------------------------------
|
| 37 |
+
fastapi==0.133.0
|
| 38 |
+
# Pinned explicitly: FastAPI's own range (starlette>=0.40.0) still admits
|
| 39 |
+
# Starlette versions with known advisories.
|
| 40 |
+
starlette==1.3.1
|
| 41 |
+
uvicorn[standard]==0.30.1
|
| 42 |
+
python-multipart==0.0.31
|
| 43 |
+
|
| 44 |
+
# ---- Model artefact pull at startup ------------------------------------------
|
| 45 |
+
huggingface-hub==0.23.4
|
| 46 |
+
|
| 47 |
+
# ---- Logging / async / CLI ---------------------------------------------------
|
| 48 |
+
structlog==24.2.0
|
| 49 |
+
anyio==4.14.2
|
| 50 |
+
tqdm==4.66.4
|
| 51 |
+
click==8.1.7
|
src/captioning/__init__.py
ADDED
|
@@ -0,0 +1,23 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Captioning — production-grade extraction of the IEEE image-captioning research.
|
| 2 |
+
|
| 3 |
+
The package mirrors the IEEE notebook
|
| 4 |
+
(``notebooks/01_ieee_inceptionv3_transformer.ipynb``) but separates orthogonal
|
| 5 |
+
concerns into sub-packages so each piece is independently testable, composable,
|
| 6 |
+
and reusable from FastAPI / scripts.
|
| 7 |
+
|
| 8 |
+
Sub-package map:
|
| 9 |
+
config/ Pydantic settings + YAML loader (the project's "type system")
|
| 10 |
+
preprocessing/ Pure transforms on captions and images (no I/O, no state)
|
| 11 |
+
data/ COCO loaders, splits, tf.data pipelines (I/O + statefulness)
|
| 12 |
+
models/ Keras layers and models (CNN encoder + Transformer decoder)
|
| 13 |
+
training/ Losses, callbacks, training orchestration
|
| 14 |
+
inference/ Generation algorithms + a singleton-friendly Predictor
|
| 15 |
+
evaluation/ BLEU/CIDEr/METEOR/ROUGE (Phase 1b expands these)
|
| 16 |
+
baselines/ Phase 3 captioner interface: CNN + Hugging Face adapters
|
| 17 |
+
utils/ Cross-cutting helpers (logging, seed, hashing, paths)
|
| 18 |
+
|
| 19 |
+
Public API is intentionally small. Everything else is internal and may change.
|
| 20 |
+
"""
|
| 21 |
+
|
| 22 |
+
__version__ = "0.1.0"
|
| 23 |
+
__all__ = ["__version__"]
|
src/captioning/baselines/__init__.py
ADDED
|
@@ -0,0 +1,24 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Phase 3 comparison models behind one captioning interface (ADR-019).
|
| 2 |
+
|
| 3 |
+
Every compared model (the project's CNN + Transformer and the pretrained
|
| 4 |
+
Hugging Face baselines) captions images through :class:`Captioner`, so the
|
| 5 |
+
comparison runner and the latency benchmark share one call path and one
|
| 6 |
+
caption normalisation.
|
| 7 |
+
|
| 8 |
+
``transformers`` and ``torch`` belong to the optional ``[hf]`` extra and are
|
| 9 |
+
imported only inside :meth:`HFCaptioner.load`. Importing this package loads
|
| 10 |
+
neither of them, and doesn't load TensorFlow either.
|
| 11 |
+
"""
|
| 12 |
+
|
| 13 |
+
from captioning.baselines.base import Captioner, CaptionerIdentity
|
| 14 |
+
from captioning.baselines.cnn import CNNCaptioner
|
| 15 |
+
from captioning.baselines.hf import HF_INSTALL_HINT, HFCaptioner, MissingHFDependencyError
|
| 16 |
+
|
| 17 |
+
__all__ = [
|
| 18 |
+
"HF_INSTALL_HINT",
|
| 19 |
+
"CNNCaptioner",
|
| 20 |
+
"Captioner",
|
| 21 |
+
"CaptionerIdentity",
|
| 22 |
+
"HFCaptioner",
|
| 23 |
+
"MissingHFDependencyError",
|
| 24 |
+
]
|
src/captioning/baselines/base.py
ADDED
|
@@ -0,0 +1,64 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""The captioner interface shared by every Phase 3 model."""
|
| 2 |
+
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
|
| 5 |
+
from abc import ABC, abstractmethod
|
| 6 |
+
from collections.abc import Mapping, Sequence
|
| 7 |
+
from dataclasses import dataclass, field
|
| 8 |
+
from pathlib import Path
|
| 9 |
+
from types import MappingProxyType
|
| 10 |
+
|
| 11 |
+
from captioning.evaluation.tokenization import strip_sentinels
|
| 12 |
+
from captioning.preprocessing.caption import preprocess_caption
|
| 13 |
+
|
| 14 |
+
|
| 15 |
+
@dataclass(frozen=True)
|
| 16 |
+
class CaptionerIdentity:
|
| 17 |
+
"""What a run records about the model behind its captions.
|
| 18 |
+
|
| 19 |
+
Attributes:
|
| 20 |
+
model_id: The ``run_meta.json`` model id (``EVAL_METHODOLOGY.md`` § 8.1).
|
| 21 |
+
hub_repo: Hugging Face Hub repository the weights come from, if any.
|
| 22 |
+
revision: Pinned commit SHA of ``hub_repo``, if any.
|
| 23 |
+
decode_settings: The decoding parameters actually used. Read-only.
|
| 24 |
+
"""
|
| 25 |
+
|
| 26 |
+
model_id: str
|
| 27 |
+
hub_repo: str | None
|
| 28 |
+
revision: str | None
|
| 29 |
+
decode_settings: Mapping[str, object] = field(default_factory=dict)
|
| 30 |
+
|
| 31 |
+
def __post_init__(self) -> None:
|
| 32 |
+
object.__setattr__(self, "decode_settings", MappingProxyType(dict(self.decode_settings)))
|
| 33 |
+
|
| 34 |
+
|
| 35 |
+
class Captioner(ABC):
|
| 36 |
+
"""Captions a batch of images and normalises the result.
|
| 37 |
+
|
| 38 |
+
Subclasses produce raw captions; :meth:`caption` normalises every one of
|
| 39 |
+
them through the existing training path (``preprocess_caption``, then
|
| 40 |
+
``strip_sentinels``), so all compared models are scored on the same
|
| 41 |
+
footing (``EVAL_METHODOLOGY.md`` § 8.3).
|
| 42 |
+
"""
|
| 43 |
+
|
| 44 |
+
def __init__(self, identity: CaptionerIdentity) -> None:
|
| 45 |
+
self.identity = identity
|
| 46 |
+
|
| 47 |
+
def load(self) -> None: # noqa: B027 - optional hook, a no-op by default
|
| 48 |
+
"""Load model weights if they aren't loaded yet. Safe to call repeatedly."""
|
| 49 |
+
|
| 50 |
+
def caption(self, image_paths: Sequence[str | Path]) -> list[str]:
|
| 51 |
+
"""Return one normalised caption per image, in input order."""
|
| 52 |
+
paths = [Path(p) for p in image_paths]
|
| 53 |
+
if not paths:
|
| 54 |
+
return []
|
| 55 |
+
raw = self._raw_captions(paths)
|
| 56 |
+
if len(raw) != len(paths):
|
| 57 |
+
raise RuntimeError(
|
| 58 |
+
f"{self.identity.model_id}: {len(raw)} captions for {len(paths)} images"
|
| 59 |
+
)
|
| 60 |
+
return [strip_sentinels(preprocess_caption(text)) for text in raw]
|
| 61 |
+
|
| 62 |
+
@abstractmethod
|
| 63 |
+
def _raw_captions(self, image_paths: list[Path]) -> list[str]:
|
| 64 |
+
"""Return the model's captions for ``image_paths``, before normalisation."""
|
src/captioning/baselines/cnn.py
ADDED
|
@@ -0,0 +1,101 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""The project's CNN + Transformer as a :class:`Captioner`."""
|
| 2 |
+
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
|
| 5 |
+
from pathlib import Path
|
| 6 |
+
from typing import TYPE_CHECKING, Any, Protocol
|
| 7 |
+
|
| 8 |
+
from captioning.baselines.base import Captioner, CaptionerIdentity
|
| 9 |
+
from captioning.config.schema import AppConfig, ComparedModelConfig
|
| 10 |
+
|
| 11 |
+
if TYPE_CHECKING:
|
| 12 |
+
from captioning.inference.predictor import DecodeStrategy
|
| 13 |
+
|
| 14 |
+
|
| 15 |
+
class _ModelSettings(Protocol):
|
| 16 |
+
max_length: int
|
| 17 |
+
|
| 18 |
+
|
| 19 |
+
class _PredictorConfig(Protocol):
|
| 20 |
+
@property
|
| 21 |
+
def model(self) -> _ModelSettings: ...
|
| 22 |
+
|
| 23 |
+
|
| 24 |
+
class PathPredictor(Protocol):
|
| 25 |
+
"""The part of ``CaptionPredictor`` the adapter relies on."""
|
| 26 |
+
|
| 27 |
+
decode_strategy: str
|
| 28 |
+
beam_width: int
|
| 29 |
+
length_penalty: float
|
| 30 |
+
repetition_penalty: float
|
| 31 |
+
no_repeat_ngram_size: int
|
| 32 |
+
|
| 33 |
+
@property
|
| 34 |
+
def config(self) -> _PredictorConfig: ...
|
| 35 |
+
|
| 36 |
+
def predict_path(self, image_path: str | Path) -> str: ...
|
| 37 |
+
|
| 38 |
+
|
| 39 |
+
class CNNCaptioner(Captioner):
|
| 40 |
+
"""Wraps a loaded ``CaptionPredictor`` without changing how it decodes.
|
| 41 |
+
|
| 42 |
+
Each image goes through ``predictor.predict_path``, one at a time, exactly
|
| 43 |
+
as ``scripts/evaluate.py`` does. The recorded decode settings follow
|
| 44 |
+
``run_meta.json``: beam-only parameters are ``None`` for greedy decoding.
|
| 45 |
+
"""
|
| 46 |
+
|
| 47 |
+
def __init__(self, predictor: PathPredictor, model: ComparedModelConfig) -> None:
|
| 48 |
+
beam = predictor.decode_strategy == "beam"
|
| 49 |
+
super().__init__(
|
| 50 |
+
CaptionerIdentity(
|
| 51 |
+
model_id=model.model_id,
|
| 52 |
+
hub_repo=model.hub_repo,
|
| 53 |
+
revision=model.revision,
|
| 54 |
+
decode_settings={
|
| 55 |
+
"decode_strategy": predictor.decode_strategy,
|
| 56 |
+
"beam_width": predictor.beam_width if beam else None,
|
| 57 |
+
"length_penalty": predictor.length_penalty if beam else None,
|
| 58 |
+
"repetition_penalty": predictor.repetition_penalty,
|
| 59 |
+
"no_repeat_ngram_size": predictor.no_repeat_ngram_size,
|
| 60 |
+
"max_length": predictor.config.model.max_length,
|
| 61 |
+
},
|
| 62 |
+
)
|
| 63 |
+
)
|
| 64 |
+
self._predictor = predictor
|
| 65 |
+
|
| 66 |
+
@classmethod
|
| 67 |
+
def from_artifacts(
|
| 68 |
+
cls,
|
| 69 |
+
weights_path: str | Path,
|
| 70 |
+
tokenizer_dir: str | Path,
|
| 71 |
+
config: AppConfig,
|
| 72 |
+
*,
|
| 73 |
+
decode_strategy: DecodeStrategy | None = None,
|
| 74 |
+
beam_width: int | None = None,
|
| 75 |
+
length_penalty: float | None = None,
|
| 76 |
+
repetition_penalty: float | None = None,
|
| 77 |
+
no_repeat_ngram_size: int | None = None,
|
| 78 |
+
) -> CNNCaptioner:
|
| 79 |
+
"""Load the checkpoint with ``CaptionPredictor.from_artifacts``.
|
| 80 |
+
|
| 81 |
+
Decoding arguments are passed through unchanged, so unset ones fall
|
| 82 |
+
back to ``config.serve`` just as they do for ``scripts/evaluate.py``.
|
| 83 |
+
The identity comes from ``config.compare.cnn``. TensorFlow is imported
|
| 84 |
+
here, not when this module is imported.
|
| 85 |
+
"""
|
| 86 |
+
from captioning.inference.predictor import CaptionPredictor
|
| 87 |
+
|
| 88 |
+
predictor: Any = CaptionPredictor.from_artifacts(
|
| 89 |
+
weights_path,
|
| 90 |
+
tokenizer_dir,
|
| 91 |
+
config,
|
| 92 |
+
decode_strategy=decode_strategy,
|
| 93 |
+
beam_width=beam_width,
|
| 94 |
+
length_penalty=length_penalty,
|
| 95 |
+
repetition_penalty=repetition_penalty,
|
| 96 |
+
no_repeat_ngram_size=no_repeat_ngram_size,
|
| 97 |
+
)
|
| 98 |
+
return cls(predictor, config.compare.cnn)
|
| 99 |
+
|
| 100 |
+
def _raw_captions(self, image_paths: list[Path]) -> list[str]:
|
| 101 |
+
return [self._predictor.predict_path(path) for path in image_paths]
|
src/captioning/baselines/hf.py
ADDED
|
@@ -0,0 +1,126 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Pretrained Hugging Face captioning baselines as a :class:`Captioner`.
|
| 2 |
+
|
| 3 |
+
``torch`` and ``transformers`` (the optional ``[hf]`` extra) are imported only
|
| 4 |
+
inside :meth:`HFCaptioner.load`, through ``importlib``, so importing this
|
| 5 |
+
module never loads them (ADR-019).
|
| 6 |
+
"""
|
| 7 |
+
|
| 8 |
+
from __future__ import annotations
|
| 9 |
+
|
| 10 |
+
import importlib
|
| 11 |
+
from dataclasses import dataclass
|
| 12 |
+
from pathlib import Path
|
| 13 |
+
from typing import Any
|
| 14 |
+
|
| 15 |
+
from PIL import Image
|
| 16 |
+
|
| 17 |
+
from captioning.baselines.base import Captioner, CaptionerIdentity
|
| 18 |
+
from captioning.config.schema import BaselineDecodeConfig, ComparedModelConfig
|
| 19 |
+
|
| 20 |
+
HF_INSTALL_HINT = 'pip install -e ".[hf]"'
|
| 21 |
+
|
| 22 |
+
|
| 23 |
+
class MissingHFDependencyError(ImportError):
|
| 24 |
+
"""A Hugging Face baseline was used without the ``[hf]`` extra installed."""
|
| 25 |
+
|
| 26 |
+
|
| 27 |
+
@dataclass(frozen=True)
|
| 28 |
+
class _Loaded:
|
| 29 |
+
torch: Any
|
| 30 |
+
dtype: Any
|
| 31 |
+
image_processor: Any
|
| 32 |
+
tokenizer: Any
|
| 33 |
+
model: Any
|
| 34 |
+
|
| 35 |
+
|
| 36 |
+
class HFCaptioner(Captioner):
|
| 37 |
+
"""One Hugging Face image-to-text model, pinned to an exact revision.
|
| 38 |
+
|
| 39 |
+
BLIP-base, ViT-GPT2 and GIT-base-coco all load through
|
| 40 |
+
``AutoImageProcessor``, ``AutoTokenizer`` and ``AutoModelForVision2Seq``.
|
| 41 |
+
Generation uses only the protocol's decode settings, passed explicitly to
|
| 42 |
+
``generate()``, with no text prompt (``EVAL_METHODOLOGY.md`` § 8.4).
|
| 43 |
+
|
| 44 |
+
Construction is cheap and needs no ``[hf]`` packages. Weights load on
|
| 45 |
+
:meth:`load`, or on the first :meth:`caption` call.
|
| 46 |
+
"""
|
| 47 |
+
|
| 48 |
+
def __init__(
|
| 49 |
+
self,
|
| 50 |
+
model: ComparedModelConfig,
|
| 51 |
+
decode: BaselineDecodeConfig,
|
| 52 |
+
*,
|
| 53 |
+
device: str = "cpu",
|
| 54 |
+
) -> None:
|
| 55 |
+
super().__init__(
|
| 56 |
+
CaptionerIdentity(
|
| 57 |
+
model_id=model.model_id,
|
| 58 |
+
hub_repo=model.hub_repo,
|
| 59 |
+
revision=model.revision,
|
| 60 |
+
decode_settings=decode.model_dump(),
|
| 61 |
+
)
|
| 62 |
+
)
|
| 63 |
+
self._model = model
|
| 64 |
+
self._decode = decode
|
| 65 |
+
self._device = device
|
| 66 |
+
self._loaded: _Loaded | None = None
|
| 67 |
+
|
| 68 |
+
@property
|
| 69 |
+
def generate_kwargs(self) -> dict[str, Any]:
|
| 70 |
+
"""The exact keyword arguments passed to ``model.generate()``."""
|
| 71 |
+
return {
|
| 72 |
+
"num_beams": self._decode.num_beams,
|
| 73 |
+
"do_sample": self._decode.do_sample,
|
| 74 |
+
"max_new_tokens": self._decode.max_new_tokens,
|
| 75 |
+
"repetition_penalty": self._decode.repetition_penalty,
|
| 76 |
+
"no_repeat_ngram_size": self._decode.no_repeat_ngram_size,
|
| 77 |
+
}
|
| 78 |
+
|
| 79 |
+
def load(self) -> None:
|
| 80 |
+
"""Load the processor, tokenizer and model at the pinned revision.
|
| 81 |
+
|
| 82 |
+
Raises:
|
| 83 |
+
MissingHFDependencyError: If ``torch`` or ``transformers`` is not
|
| 84 |
+
installed.
|
| 85 |
+
"""
|
| 86 |
+
if self._loaded is not None:
|
| 87 |
+
return
|
| 88 |
+
torch, transformers = _import_hf()
|
| 89 |
+
repo, revision = self._model.hub_repo, self._model.revision
|
| 90 |
+
dtype = getattr(torch, self._decode.precision)
|
| 91 |
+
image_processor = transformers.AutoImageProcessor.from_pretrained(repo, revision=revision)
|
| 92 |
+
tokenizer = transformers.AutoTokenizer.from_pretrained(repo, revision=revision)
|
| 93 |
+
model = transformers.AutoModelForVision2Seq.from_pretrained(
|
| 94 |
+
repo, revision=revision, torch_dtype=dtype
|
| 95 |
+
)
|
| 96 |
+
model = model.to(self._device)
|
| 97 |
+
model.eval()
|
| 98 |
+
self._loaded = _Loaded(torch, dtype, image_processor, tokenizer, model)
|
| 99 |
+
|
| 100 |
+
def _raw_captions(self, image_paths: list[Path]) -> list[str]:
|
| 101 |
+
self.load()
|
| 102 |
+
loaded = self._loaded
|
| 103 |
+
assert loaded is not None
|
| 104 |
+
images = [_open_rgb(path) for path in image_paths]
|
| 105 |
+
pixel_values = loaded.image_processor(images=images, return_tensors="pt").pixel_values
|
| 106 |
+
pixel_values = pixel_values.to(self._device, dtype=loaded.dtype)
|
| 107 |
+
with loaded.torch.inference_mode():
|
| 108 |
+
output_ids = loaded.model.generate(pixel_values=pixel_values, **self.generate_kwargs)
|
| 109 |
+
return list(loaded.tokenizer.batch_decode(output_ids, skip_special_tokens=True))
|
| 110 |
+
|
| 111 |
+
|
| 112 |
+
def _import_hf() -> tuple[Any, Any]:
|
| 113 |
+
try:
|
| 114 |
+
torch = importlib.import_module("torch")
|
| 115 |
+
transformers = importlib.import_module("transformers")
|
| 116 |
+
except ImportError as exc:
|
| 117 |
+
raise MissingHFDependencyError(
|
| 118 |
+
"Hugging Face baselines need the optional [hf] extra (transformers, torch). "
|
| 119 |
+
f"Install it with: {HF_INSTALL_HINT}"
|
| 120 |
+
) from exc
|
| 121 |
+
return torch, transformers
|
| 122 |
+
|
| 123 |
+
|
| 124 |
+
def _open_rgb(path: Path) -> Image.Image:
|
| 125 |
+
with Image.open(path) as image:
|
| 126 |
+
return image.convert("RGB")
|
src/captioning/config/__init__.py
ADDED
|
@@ -0,0 +1,30 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Configuration package — Pydantic schemas and YAML loaders.
|
| 2 |
+
|
| 3 |
+
Why a dedicated package? Configs are the project's *type system*. Every other
|
| 4 |
+
module accepts an `AppConfig` (or a sub-config) instead of pulling globals,
|
| 5 |
+
which makes them testable in isolation and trivially overridable in CI / serve.
|
| 6 |
+
"""
|
| 7 |
+
|
| 8 |
+
from captioning.config.loader import load_config
|
| 9 |
+
from captioning.config.schema import (
|
| 10 |
+
AppConfig,
|
| 11 |
+
BaselineDecodeConfig,
|
| 12 |
+
CompareConfig,
|
| 13 |
+
ComparedModelConfig,
|
| 14 |
+
DataConfig,
|
| 15 |
+
ModelConfig,
|
| 16 |
+
ServeConfig,
|
| 17 |
+
TrainConfig,
|
| 18 |
+
)
|
| 19 |
+
|
| 20 |
+
__all__ = [
|
| 21 |
+
"AppConfig",
|
| 22 |
+
"BaselineDecodeConfig",
|
| 23 |
+
"CompareConfig",
|
| 24 |
+
"ComparedModelConfig",
|
| 25 |
+
"DataConfig",
|
| 26 |
+
"ModelConfig",
|
| 27 |
+
"ServeConfig",
|
| 28 |
+
"TrainConfig",
|
| 29 |
+
"load_config",
|
| 30 |
+
]
|
src/captioning/config/loader.py
ADDED
|
@@ -0,0 +1,73 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""YAML-to-Pydantic config loader.
|
| 2 |
+
|
| 3 |
+
Why this exists separately from ``schema.py``:
|
| 4 |
+
* Schema is *what* a valid config looks like; loader is *how* you build one.
|
| 5 |
+
Splitting them lets tests build an ``AppConfig`` programmatically without
|
| 6 |
+
touching disk, and lets the loader gain features (env-file resolution,
|
| 7 |
+
multi-file merging) without changing the schema.
|
| 8 |
+
"""
|
| 9 |
+
|
| 10 |
+
from __future__ import annotations
|
| 11 |
+
|
| 12 |
+
from pathlib import Path
|
| 13 |
+
from typing import Any
|
| 14 |
+
|
| 15 |
+
import yaml
|
| 16 |
+
from pydantic_settings import EnvSettingsSource
|
| 17 |
+
|
| 18 |
+
from captioning.config.schema import AppConfig
|
| 19 |
+
|
| 20 |
+
|
| 21 |
+
def load_config(path: str | Path) -> AppConfig:
|
| 22 |
+
"""Load a YAML file into an ``AppConfig``, apply env overrides, and validate it.
|
| 23 |
+
|
| 24 |
+
Precedence, highest first: ``CAPTIONING__*`` environment variables, the
|
| 25 |
+
YAML file, the schema defaults. An override replaces only the field it
|
| 26 |
+
names (``CAPTIONING__SERVE__BEAM_WIDTH=4`` keeps the rest of ``serve``);
|
| 27 |
+
a list value, such as ``CAPTIONING__SERVE__CORS_ALLOWED_ORIGINS``,
|
| 28 |
+
replaces the YAML list whole.
|
| 29 |
+
|
| 30 |
+
Args:
|
| 31 |
+
path: Path to a YAML file with the structure::
|
| 32 |
+
|
| 33 |
+
data: {...}
|
| 34 |
+
model: {...}
|
| 35 |
+
train: {...}
|
| 36 |
+
serve: {...}
|
| 37 |
+
|
| 38 |
+
Returns:
|
| 39 |
+
A fully validated, immutable ``AppConfig`` instance.
|
| 40 |
+
|
| 41 |
+
Raises:
|
| 42 |
+
FileNotFoundError: If the YAML path does not exist.
|
| 43 |
+
pydantic.ValidationError: If any field fails validation.
|
| 44 |
+
pydantic_settings.sources.SettingsError: If an env value for a list or
|
| 45 |
+
section isn't valid JSON.
|
| 46 |
+
"""
|
| 47 |
+
path = Path(path)
|
| 48 |
+
if not path.is_file():
|
| 49 |
+
raise FileNotFoundError(f"Config file not found: {path}")
|
| 50 |
+
|
| 51 |
+
with path.open(encoding="utf-8") as f:
|
| 52 |
+
raw: dict[str, Any] = yaml.safe_load(f) or {}
|
| 53 |
+
|
| 54 |
+
# pydantic-settings ranks constructor arguments above the environment, so
|
| 55 |
+
# ``AppConfig(**raw)`` alone would let the YAML win. Merge the env values
|
| 56 |
+
# over it first, parsed by the same source ``AppConfig`` uses.
|
| 57 |
+
env = EnvSettingsSource(AppConfig)()
|
| 58 |
+
return AppConfig(**_merge(raw, env))
|
| 59 |
+
|
| 60 |
+
|
| 61 |
+
def _merge(base: dict[str, Any], override: dict[str, Any]) -> dict[str, Any]:
|
| 62 |
+
"""Return ``base`` updated by ``override``, recursing into sections both set.
|
| 63 |
+
|
| 64 |
+
Any other value, lists included, is replaced whole, as pydantic-settings
|
| 65 |
+
merges its own sources.
|
| 66 |
+
"""
|
| 67 |
+
merged = dict(base)
|
| 68 |
+
for key, value in override.items():
|
| 69 |
+
if isinstance(value, dict) and isinstance(merged.get(key), dict):
|
| 70 |
+
merged[key] = _merge(merged[key], value)
|
| 71 |
+
else:
|
| 72 |
+
merged[key] = value
|
| 73 |
+
return merged
|
src/captioning/config/schema.py
ADDED
|
@@ -0,0 +1,289 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Typed configuration schemas (Pydantic v2 ``BaseSettings``).
|
| 2 |
+
|
| 3 |
+
These classes replace the bare globals ``MAX_LENGTH``, ``BATCH_SIZE``, ... that
|
| 4 |
+
the notebook holds in cell 6. The advantages of doing this:
|
| 5 |
+
|
| 6 |
+
1. **Type safety** — every field has a declared type and Pydantic validates
|
| 7 |
+
it at load time. A YAML typo (``batch_size: "64"`` as a string) raises an
|
| 8 |
+
error pointing at the file and field, not a mysterious training failure
|
| 9 |
+
six steps later.
|
| 10 |
+
2. **Env override** — ``CAPTIONING__TRAIN__BATCH_SIZE=32`` overrides
|
| 11 |
+
``train.batch_size`` without editing YAML. The double underscore is the
|
| 12 |
+
nesting delimiter (configurable below). Useful for CI smoke tests.
|
| 13 |
+
3. **Single source of truth** — every other module accepts a sub-config
|
| 14 |
+
(``ModelConfig``, ``TrainConfig``, ...) instead of pulling globals. That
|
| 15 |
+
makes them testable in isolation and trivially overridable in serve.
|
| 16 |
+
|
| 17 |
+
The schema mirrors the IEEE notebook 1:1 — same field names where reasonable,
|
| 18 |
+
same default values. Extending it (Phase 1b: warmup/cosine LR; Phase 3: model
|
| 19 |
+
registry) only adds new fields, never changes the meaning of existing ones.
|
| 20 |
+
"""
|
| 21 |
+
|
| 22 |
+
from __future__ import annotations
|
| 23 |
+
|
| 24 |
+
from pathlib import Path
|
| 25 |
+
from typing import Literal
|
| 26 |
+
|
| 27 |
+
from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
|
| 28 |
+
from pydantic_settings import BaseSettings, SettingsConfigDict
|
| 29 |
+
|
| 30 |
+
|
| 31 |
+
class _StrictModel(BaseModel):
|
| 32 |
+
"""Shared base for every sub-config — rejects unknown keys.
|
| 33 |
+
|
| 34 |
+
Pydantic's default ``extra="ignore"`` silently drops misspelled fields.
|
| 35 |
+
For configs that drive ML hyperparameters that's the worst possible
|
| 36 |
+
behaviour: a typo (``vocabularsy_size`` instead of ``vocabulary_size``)
|
| 37 |
+
silently uses the default and the model trains with the wrong value.
|
| 38 |
+
Forbidding extras turns every typo into a load-time error pointing at
|
| 39 |
+
the offending field.
|
| 40 |
+
|
| 41 |
+
Note: ``extra="forbid"`` is set on ``AppConfig`` separately because
|
| 42 |
+
``BaseSettings`` uses ``SettingsConfigDict``, not ``ConfigDict``.
|
| 43 |
+
"""
|
| 44 |
+
|
| 45 |
+
model_config = ConfigDict(extra="forbid")
|
| 46 |
+
|
| 47 |
+
|
| 48 |
+
class DataConfig(_StrictModel):
|
| 49 |
+
"""Where the dataset lives and how much of it to use.
|
| 50 |
+
|
| 51 |
+
Attributes:
|
| 52 |
+
base_path: Root of the COCO dataset. Mirrors the notebook's
|
| 53 |
+
``BASE_PATH = '../input/coco-2017-dataset/coco2017'``.
|
| 54 |
+
annotations_filename: Name of the captions JSON inside ``annotations/``.
|
| 55 |
+
images_subdir: Sub-folder under ``base_path`` containing JPEGs.
|
| 56 |
+
sample_size: How many caption pairs to sample. The notebook samples
|
| 57 |
+
120k. Set to ``-1`` to use the full set.
|
| 58 |
+
train_val_split: Fraction of *images* (not captions) used for training.
|
| 59 |
+
Splitting at the image level prevents the same image appearing in
|
| 60 |
+
both splits via different captions — a real leakage source.
|
| 61 |
+
"""
|
| 62 |
+
|
| 63 |
+
base_path: Path = Path("data/coco2017")
|
| 64 |
+
annotations_filename: str = "captions_train2017.json"
|
| 65 |
+
images_subdir: str = "train2017"
|
| 66 |
+
sample_size: int = 120_000
|
| 67 |
+
train_val_split: float = 0.8
|
| 68 |
+
|
| 69 |
+
@field_validator("train_val_split")
|
| 70 |
+
@classmethod
|
| 71 |
+
def _validate_split(cls, v: float) -> float:
|
| 72 |
+
if not 0.0 < v < 1.0:
|
| 73 |
+
raise ValueError(f"train_val_split must be in (0, 1), got {v}")
|
| 74 |
+
return v
|
| 75 |
+
|
| 76 |
+
|
| 77 |
+
class ModelConfig(_StrictModel):
|
| 78 |
+
"""Architecture hyperparameters.
|
| 79 |
+
|
| 80 |
+
Defaults match the IEEE paper / notebook cell 6 exactly. Changing any of
|
| 81 |
+
these requires re-training and re-publishing the model card on HF Hub.
|
| 82 |
+
"""
|
| 83 |
+
|
| 84 |
+
embedding_dim: int = 512
|
| 85 |
+
units: int = 512
|
| 86 |
+
max_length: int = 40
|
| 87 |
+
vocabulary_size: int = 15_000
|
| 88 |
+
encoder_num_heads: int = 1 # Notebook cell 21: TransformerEncoderLayer(EMBEDDING_DIM, 1)
|
| 89 |
+
decoder_num_heads: int = 8 # Notebook cell 21: TransformerDecoderLayer(..., 8)
|
| 90 |
+
decoder_dropout_inner: float = 0.3 # Notebook cell 19: dropout_1
|
| 91 |
+
decoder_dropout_outer: float = 0.5 # Notebook cell 19: dropout_2
|
| 92 |
+
decoder_attention_dropout: float = 0.1 # Notebook cell 19: MultiHeadAttention(dropout=0.1)
|
| 93 |
+
|
| 94 |
+
|
| 95 |
+
class TrainConfig(_StrictModel):
|
| 96 |
+
"""Optimisation hyperparameters.
|
| 97 |
+
|
| 98 |
+
The Phase 1 baseline mirrors the IEEE notebook: constant LR, no label
|
| 99 |
+
smoothing, dropout-active validation (a notebook quirk preserved for
|
| 100 |
+
parity). The fields below the comment line are *opt-in* training-
|
| 101 |
+
stability knobs added during the caption-quality stabilisation phase.
|
| 102 |
+
Defaults keep every existing run byte-for-byte identical to the
|
| 103 |
+
notebook; flipping the flag in YAML opts a run into the modern recipe.
|
| 104 |
+
"""
|
| 105 |
+
|
| 106 |
+
epochs: int = 10
|
| 107 |
+
batch_size: int = 64
|
| 108 |
+
buffer_size: int = 1_000 # tf.data shuffle buffer
|
| 109 |
+
early_stopping_patience: int = 3
|
| 110 |
+
seed: int = 42 # NEW (not in notebook): pin RNGs for reproducibility
|
| 111 |
+
learning_rate: float = 1e-3 # Notebook uses Keras Adam default == 1e-3
|
| 112 |
+
weights_filename: str = "model.h5"
|
| 113 |
+
|
| 114 |
+
# ---- opt-in stability flags (default values preserve notebook parity) ----
|
| 115 |
+
label_smoothing: float = 0.0
|
| 116 |
+
lr_schedule: str = "constant" # "constant" | "cosine"
|
| 117 |
+
warmup_steps: int = 0
|
| 118 |
+
cosine_decay_steps: int | None = None # If None, derived from epochs * steps_per_epoch
|
| 119 |
+
min_learning_rate: float = 0.0
|
| 120 |
+
honour_training_flag_in_test_step: bool = False # parity-quirk override
|
| 121 |
+
|
| 122 |
+
@field_validator("label_smoothing")
|
| 123 |
+
@classmethod
|
| 124 |
+
def _validate_label_smoothing(cls, v: float) -> float:
|
| 125 |
+
if not 0.0 <= v < 1.0:
|
| 126 |
+
raise ValueError(f"label_smoothing must be in [0, 1), got {v}")
|
| 127 |
+
return v
|
| 128 |
+
|
| 129 |
+
@field_validator("lr_schedule")
|
| 130 |
+
@classmethod
|
| 131 |
+
def _validate_lr_schedule(cls, v: str) -> str:
|
| 132 |
+
if v not in {"constant", "cosine"}:
|
| 133 |
+
raise ValueError(f"lr_schedule must be 'constant' or 'cosine', got {v!r}")
|
| 134 |
+
return v
|
| 135 |
+
|
| 136 |
+
@field_validator("warmup_steps")
|
| 137 |
+
@classmethod
|
| 138 |
+
def _validate_warmup_steps(cls, v: int) -> int:
|
| 139 |
+
if v < 0:
|
| 140 |
+
raise ValueError(f"warmup_steps must be >= 0, got {v}")
|
| 141 |
+
return v
|
| 142 |
+
|
| 143 |
+
|
| 144 |
+
class ServeConfig(_StrictModel):
|
| 145 |
+
"""Settings for the FastAPI backend (Phase 2). Defined here so the schema
|
| 146 |
+
is complete and tests don't have to mock a sub-config's existence.
|
| 147 |
+
|
| 148 |
+
Decoding-related defaults are deliberately conservative: ``greedy`` stays
|
| 149 |
+
the default for byte-for-byte parity with the IEEE notebook. Switching to
|
| 150 |
+
beam at deploy time is a one-line YAML override:
|
| 151 |
+
|
| 152 |
+
serve:
|
| 153 |
+
decode_strategy: beam
|
| 154 |
+
beam_width: 4
|
| 155 |
+
length_penalty: 0.7
|
| 156 |
+
repetition_penalty: 1.1
|
| 157 |
+
no_repeat_ngram_size: 3
|
| 158 |
+
"""
|
| 159 |
+
|
| 160 |
+
max_upload_bytes: int = 10 * 1024 * 1024 # 10 MB
|
| 161 |
+
decode_strategy: str = "greedy"
|
| 162 |
+
beam_width: int = 3
|
| 163 |
+
length_penalty: float = 1.0
|
| 164 |
+
repetition_penalty: float = 1.0
|
| 165 |
+
no_repeat_ngram_size: int = 0
|
| 166 |
+
cors_allowed_origins: list[str] = Field(default_factory=lambda: ["http://localhost:3000"])
|
| 167 |
+
|
| 168 |
+
@field_validator("decode_strategy")
|
| 169 |
+
@classmethod
|
| 170 |
+
def _validate_decode_strategy(cls, v: str) -> str:
|
| 171 |
+
if v not in {"greedy", "beam"}:
|
| 172 |
+
raise ValueError(f"decode_strategy must be 'greedy' or 'beam', got {v!r}")
|
| 173 |
+
return v
|
| 174 |
+
|
| 175 |
+
@field_validator("beam_width")
|
| 176 |
+
@classmethod
|
| 177 |
+
def _validate_beam_width(cls, v: int) -> int:
|
| 178 |
+
if v < 1:
|
| 179 |
+
raise ValueError(f"beam_width must be >= 1, got {v}")
|
| 180 |
+
return v
|
| 181 |
+
|
| 182 |
+
@field_validator("repetition_penalty")
|
| 183 |
+
@classmethod
|
| 184 |
+
def _validate_repetition_penalty(cls, v: float) -> float:
|
| 185 |
+
if v < 1.0:
|
| 186 |
+
raise ValueError(f"repetition_penalty must be >= 1.0 (1.0 disables it), got {v}")
|
| 187 |
+
return v
|
| 188 |
+
|
| 189 |
+
|
| 190 |
+
class ComparedModelConfig(_StrictModel):
|
| 191 |
+
"""One model in the Phase 3 comparison, pinned to an exact Hub commit.
|
| 192 |
+
|
| 193 |
+
``revision`` must be a full 40-character commit SHA, never a branch or tag
|
| 194 |
+
name, so a moved branch can't silently change the model being evaluated.
|
| 195 |
+
"""
|
| 196 |
+
|
| 197 |
+
# ``model_id`` matches the run_meta.json field; opt out of pydantic's
|
| 198 |
+
# protected ``model_`` namespace warning, as the API response schemas do.
|
| 199 |
+
model_config = ConfigDict(extra="forbid", protected_namespaces=())
|
| 200 |
+
|
| 201 |
+
model_id: str = Field(min_length=1)
|
| 202 |
+
hub_repo: str = Field(pattern=r"^[\w.-]+/[\w.-]+$")
|
| 203 |
+
revision: str = Field(pattern=r"^[0-9a-f]{40}$")
|
| 204 |
+
|
| 205 |
+
|
| 206 |
+
class BaselineDecodeConfig(_StrictModel):
|
| 207 |
+
"""Generation settings passed explicitly to every Hugging Face baseline.
|
| 208 |
+
|
| 209 |
+
Greedy decoding, so each repository's own generation defaults never apply
|
| 210 |
+
(``docs/EVAL_METHODOLOGY.md`` § 8.4).
|
| 211 |
+
"""
|
| 212 |
+
|
| 213 |
+
num_beams: int = Field(default=1, ge=1)
|
| 214 |
+
do_sample: bool = False
|
| 215 |
+
max_new_tokens: int = Field(default=40, ge=1)
|
| 216 |
+
repetition_penalty: float = Field(default=1.0, ge=1.0)
|
| 217 |
+
no_repeat_ngram_size: int = Field(default=0, ge=0)
|
| 218 |
+
precision: Literal["float32"] = "float32"
|
| 219 |
+
|
| 220 |
+
|
| 221 |
+
def _protocol_cnn() -> ComparedModelConfig:
|
| 222 |
+
return ComparedModelConfig(
|
| 223 |
+
model_id="inceptionv3-transformer-stabilized",
|
| 224 |
+
hub_repo="apoorvrajdev/captioning-inceptionv3-transformer",
|
| 225 |
+
revision="59d93b4babb16b0ac81eef598f3abc271a355cbf", # tag v2.0.0
|
| 226 |
+
)
|
| 227 |
+
|
| 228 |
+
|
| 229 |
+
def _protocol_baselines() -> list[ComparedModelConfig]:
|
| 230 |
+
return [
|
| 231 |
+
ComparedModelConfig(
|
| 232 |
+
model_id="blip-base",
|
| 233 |
+
hub_repo="Salesforce/blip-image-captioning-base",
|
| 234 |
+
revision="82a37760796d32b1411fe092ab5d4e227313294b",
|
| 235 |
+
),
|
| 236 |
+
ComparedModelConfig(
|
| 237 |
+
model_id="vit-gpt2",
|
| 238 |
+
hub_repo="nlpconnect/vit-gpt2-image-captioning",
|
| 239 |
+
revision="dc68f91c06a1ba6f15268e5b9c13ae7a7c514084",
|
| 240 |
+
),
|
| 241 |
+
ComparedModelConfig(
|
| 242 |
+
model_id="git-base-coco",
|
| 243 |
+
hub_repo="microsoft/git-base-coco",
|
| 244 |
+
revision="a13141da42abd4a8cbf283601a8104265f537cee",
|
| 245 |
+
),
|
| 246 |
+
]
|
| 247 |
+
|
| 248 |
+
|
| 249 |
+
class CompareConfig(_StrictModel):
|
| 250 |
+
"""Phase 3 model comparison (ADR-019).
|
| 251 |
+
|
| 252 |
+
The defaults are the frozen protocol in ``docs/EVAL_METHODOLOGY.md`` § 8.1
|
| 253 |
+
(models and revisions) and § 8.4 (baseline decoding). Changing them is a
|
| 254 |
+
protocol amendment, not a tuning knob. The CNN + Transformer's own decoding
|
| 255 |
+
still comes from ``serve`` / the caller, exactly as in ``scripts/evaluate.py``.
|
| 256 |
+
"""
|
| 257 |
+
|
| 258 |
+
cnn: ComparedModelConfig = Field(default_factory=_protocol_cnn)
|
| 259 |
+
baselines: list[ComparedModelConfig] = Field(default_factory=_protocol_baselines)
|
| 260 |
+
baseline_decode: BaselineDecodeConfig = Field(default_factory=BaselineDecodeConfig)
|
| 261 |
+
|
| 262 |
+
@model_validator(mode="after")
|
| 263 |
+
def _validate_unique_model_ids(self) -> CompareConfig:
|
| 264 |
+
ids = [self.cnn.model_id, *(m.model_id for m in self.baselines)]
|
| 265 |
+
duplicates = sorted({i for i in ids if ids.count(i) > 1})
|
| 266 |
+
if duplicates:
|
| 267 |
+
raise ValueError(f"model_id must be unique across compared models, got {duplicates}")
|
| 268 |
+
return self
|
| 269 |
+
|
| 270 |
+
|
| 271 |
+
class AppConfig(BaseSettings):
|
| 272 |
+
"""Top-level config aggregating every sub-config.
|
| 273 |
+
|
| 274 |
+
Loaded by ``captioning.config.loader.load_config(yaml_path)``. Env vars
|
| 275 |
+
with prefix ``CAPTIONING__`` override fields at any depth.
|
| 276 |
+
"""
|
| 277 |
+
|
| 278 |
+
data: DataConfig = Field(default_factory=DataConfig)
|
| 279 |
+
model: ModelConfig = Field(default_factory=ModelConfig)
|
| 280 |
+
train: TrainConfig = Field(default_factory=TrainConfig)
|
| 281 |
+
serve: ServeConfig = Field(default_factory=ServeConfig)
|
| 282 |
+
compare: CompareConfig = Field(default_factory=CompareConfig)
|
| 283 |
+
|
| 284 |
+
model_config = SettingsConfigDict(
|
| 285 |
+
env_prefix="CAPTIONING__",
|
| 286 |
+
env_nested_delimiter="__",
|
| 287 |
+
case_sensitive=False,
|
| 288 |
+
extra="forbid", # Reject unknown keys — catches typos at load time
|
| 289 |
+
)
|
src/captioning/data/__init__.py
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Data — dataset loaders, splits, and tf.data pipelines.
|
| 2 |
+
|
| 3 |
+
Data lives separately from preprocessing because *I/O and randomness* are
|
| 4 |
+
fundamentally different concerns from pure transforms. This package owns:
|
| 5 |
+
|
| 6 |
+
coco.py Read COCO annotation JSONs into a (image_path, caption) DataFrame
|
| 7 |
+
splits.py Deterministic image-level train/val splitting (NOT caption-level —
|
| 8 |
+
preventing the same image from appearing in both splits)
|
| 9 |
+
pipeline.py Compose preprocessing + tokenization into tf.data pipelines
|
| 10 |
+
"""
|
| 11 |
+
|
| 12 |
+
from captioning.data.coco import load_coco_annotations
|
| 13 |
+
from captioning.data.pipeline import build_train_pipeline, build_val_pipeline
|
| 14 |
+
from captioning.data.splits import make_image_level_splits
|
| 15 |
+
|
| 16 |
+
__all__ = [
|
| 17 |
+
"build_train_pipeline",
|
| 18 |
+
"build_val_pipeline",
|
| 19 |
+
"load_coco_annotations",
|
| 20 |
+
"make_image_level_splits",
|
| 21 |
+
]
|
src/captioning/data/coco.py
ADDED
|
@@ -0,0 +1,92 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Load COCO 2017 caption annotations into a ``pandas.DataFrame``.
|
| 2 |
+
|
| 3 |
+
Mirrors notebook cell 2 with two small but important upgrades that don't
|
| 4 |
+
change behaviour at fixed seeds:
|
| 5 |
+
|
| 6 |
+
1. **Seeded sampling.** The notebook calls ``captions.sample(120000)`` with no
|
| 7 |
+
``random_state``, so two runs produce different subsets. We thread the
|
| 8 |
+
seed through so every run is identical when the seed is fixed.
|
| 9 |
+
2. **Path validation.** The notebook constructs paths via ``f-string`` with no
|
| 10 |
+
``os.path.exists`` check; if the dataset is missing, training fails ten
|
| 11 |
+
minutes in. We check the annotations file up front and raise a clear error.
|
| 12 |
+
"""
|
| 13 |
+
|
| 14 |
+
from __future__ import annotations
|
| 15 |
+
|
| 16 |
+
import json
|
| 17 |
+
from collections.abc import Callable
|
| 18 |
+
from pathlib import Path
|
| 19 |
+
|
| 20 |
+
import pandas as pd
|
| 21 |
+
|
| 22 |
+
from captioning.utils.logging import get_logger
|
| 23 |
+
|
| 24 |
+
log = get_logger(__name__)
|
| 25 |
+
|
| 26 |
+
IMAGE_FILENAME_TEMPLATE = "%012d.jpg" # Notebook cell 2: '%012d.jpg' % image_id
|
| 27 |
+
|
| 28 |
+
|
| 29 |
+
def load_coco_annotations(
|
| 30 |
+
base_path: str | Path,
|
| 31 |
+
annotations_filename: str = "captions_train2017.json",
|
| 32 |
+
images_subdir: str = "train2017",
|
| 33 |
+
sample_size: int = 120_000,
|
| 34 |
+
seed: int | None = None,
|
| 35 |
+
caption_preprocessor: Callable[[str], str] | None = None,
|
| 36 |
+
) -> pd.DataFrame:
|
| 37 |
+
"""Read COCO annotations and return a (image, caption) DataFrame.
|
| 38 |
+
|
| 39 |
+
Mirrors notebook cell 2 + cell 4 (when ``caption_preprocessor`` is
|
| 40 |
+
supplied). Returns the same columns and dtypes the notebook produces.
|
| 41 |
+
|
| 42 |
+
Args:
|
| 43 |
+
base_path: Path to the COCO root containing ``annotations/`` and the
|
| 44 |
+
images sub-directory.
|
| 45 |
+
annotations_filename: JSON file under ``base_path / 'annotations'``.
|
| 46 |
+
images_subdir: Folder containing the JPEG files.
|
| 47 |
+
sample_size: Number of caption rows to keep after sampling. Use
|
| 48 |
+
``-1`` to disable sampling and keep everything.
|
| 49 |
+
seed: Random seed for deterministic sampling. ``None`` matches the
|
| 50 |
+
notebook's non-deterministic behaviour.
|
| 51 |
+
caption_preprocessor: Optional function applied to the ``caption``
|
| 52 |
+
column. Pass ``preprocessing.preprocess_caption`` to reproduce
|
| 53 |
+
cell 4. Left optional so callers can stage the preprocessing
|
| 54 |
+
differently (the parity audit script applies it manually).
|
| 55 |
+
|
| 56 |
+
Returns:
|
| 57 |
+
``pd.DataFrame`` with columns ``image`` (absolute path) and ``caption``
|
| 58 |
+
(string), index reset.
|
| 59 |
+
|
| 60 |
+
Raises:
|
| 61 |
+
FileNotFoundError: If ``annotations_filename`` is missing.
|
| 62 |
+
"""
|
| 63 |
+
base_path = Path(base_path)
|
| 64 |
+
annotations_path = base_path / "annotations" / annotations_filename
|
| 65 |
+
if not annotations_path.is_file():
|
| 66 |
+
raise FileNotFoundError(
|
| 67 |
+
f"COCO annotations not found at {annotations_path}. "
|
| 68 |
+
f"Run `python -m scripts.prepare_data` to download."
|
| 69 |
+
)
|
| 70 |
+
|
| 71 |
+
log.info("loading_coco", path=str(annotations_path))
|
| 72 |
+
with annotations_path.open(encoding="utf-8") as f:
|
| 73 |
+
payload = json.load(f)
|
| 74 |
+
annotations = payload["annotations"]
|
| 75 |
+
|
| 76 |
+
img_dir = base_path / images_subdir
|
| 77 |
+
img_cap_pairs = [
|
| 78 |
+
[str(img_dir / (IMAGE_FILENAME_TEMPLATE % sample["image_id"])), sample["caption"]]
|
| 79 |
+
for sample in annotations
|
| 80 |
+
]
|
| 81 |
+
df = pd.DataFrame(img_cap_pairs, columns=["image", "caption"])
|
| 82 |
+
|
| 83 |
+
if sample_size > 0 and sample_size < len(df):
|
| 84 |
+
df = df.sample(n=sample_size, random_state=seed)
|
| 85 |
+
|
| 86 |
+
df = df.reset_index(drop=True)
|
| 87 |
+
|
| 88 |
+
if caption_preprocessor is not None:
|
| 89 |
+
df["caption"] = df["caption"].apply(caption_preprocessor)
|
| 90 |
+
|
| 91 |
+
log.info("coco_loaded", rows=len(df), unique_images=df["image"].nunique())
|
| 92 |
+
return df
|
src/captioning/data/pipeline.py
ADDED
|
@@ -0,0 +1,92 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""``tf.data`` pipeline construction.
|
| 2 |
+
|
| 3 |
+
Mirrors notebook cell 13 (the ``load_data`` map function) and cell 14 (the
|
| 4 |
+
two pipeline definitions). The notebook closes over a global ``tokenizer``;
|
| 5 |
+
we pass the tokenizer in explicitly so the same code works in tests, scripts,
|
| 6 |
+
and the parity audit.
|
| 7 |
+
|
| 8 |
+
Note on ``shuffle`` in the val pipeline: the notebook shuffles both train
|
| 9 |
+
and val. That is technically unnecessary for validation but harmless, and we
|
| 10 |
+
preserve it for parity. Phase 1b removes it from val.
|
| 11 |
+
"""
|
| 12 |
+
|
| 13 |
+
from __future__ import annotations
|
| 14 |
+
|
| 15 |
+
from collections.abc import Sequence
|
| 16 |
+
|
| 17 |
+
from captioning.preprocessing.image import preprocess_image_tensor
|
| 18 |
+
from captioning.preprocessing.tokenizer import CaptionTokenizer
|
| 19 |
+
|
| 20 |
+
|
| 21 |
+
def _make_load_data_fn(tokenizer: CaptionTokenizer):
|
| 22 |
+
"""Return a ``tf.data``-compatible map function (image_path, caption) -> (image, ids).
|
| 23 |
+
|
| 24 |
+
Defined as a closure rather than a top-level function so it captures the
|
| 25 |
+
tokenizer without leaking it into the module namespace. ``tf.data`` calls
|
| 26 |
+
this for every example with both arguments as ``tf.string`` tensors.
|
| 27 |
+
"""
|
| 28 |
+
import tensorflow as tf
|
| 29 |
+
|
| 30 |
+
def load_data(image_path, caption):
|
| 31 |
+
raw = tf.io.read_file(image_path)
|
| 32 |
+
image = tf.io.decode_jpeg(raw, channels=3)
|
| 33 |
+
image = preprocess_image_tensor(image)
|
| 34 |
+
ids = tokenizer.encode(caption)
|
| 35 |
+
return image, ids
|
| 36 |
+
|
| 37 |
+
return load_data
|
| 38 |
+
|
| 39 |
+
|
| 40 |
+
def build_train_pipeline(
|
| 41 |
+
image_paths: Sequence[str],
|
| 42 |
+
captions: Sequence[str],
|
| 43 |
+
tokenizer: CaptionTokenizer,
|
| 44 |
+
batch_size: int,
|
| 45 |
+
buffer_size: int,
|
| 46 |
+
):
|
| 47 |
+
"""Build the training dataset, byte-identical to notebook cell 14.
|
| 48 |
+
|
| 49 |
+
Args:
|
| 50 |
+
image_paths: One path per (image, caption) pair (image-level split
|
| 51 |
+
already applied — see ``data.splits``).
|
| 52 |
+
captions: Preprocessed captions, one per ``image_paths`` entry.
|
| 53 |
+
tokenizer: Fitted ``CaptionTokenizer``.
|
| 54 |
+
batch_size: Mini-batch size; matches ``BATCH_SIZE`` in the notebook.
|
| 55 |
+
buffer_size: Shuffle buffer size; matches ``BUFFER_SIZE``.
|
| 56 |
+
|
| 57 |
+
Returns:
|
| 58 |
+
A ``tf.data.Dataset`` yielding ``(image, token_ids)`` batches.
|
| 59 |
+
"""
|
| 60 |
+
import tensorflow as tf
|
| 61 |
+
|
| 62 |
+
load_data = _make_load_data_fn(tokenizer)
|
| 63 |
+
return (
|
| 64 |
+
tf.data.Dataset.from_tensor_slices((list(image_paths), list(captions)))
|
| 65 |
+
.map(load_data, num_parallel_calls=tf.data.AUTOTUNE)
|
| 66 |
+
.shuffle(buffer_size)
|
| 67 |
+
.batch(batch_size)
|
| 68 |
+
)
|
| 69 |
+
|
| 70 |
+
|
| 71 |
+
def build_val_pipeline(
|
| 72 |
+
image_paths: Sequence[str],
|
| 73 |
+
captions: Sequence[str],
|
| 74 |
+
tokenizer: CaptionTokenizer,
|
| 75 |
+
batch_size: int,
|
| 76 |
+
buffer_size: int,
|
| 77 |
+
):
|
| 78 |
+
"""Build the validation dataset.
|
| 79 |
+
|
| 80 |
+
Identical structure to ``build_train_pipeline``, with a separate function
|
| 81 |
+
so Phase 1b can drop the (unnecessary) shuffle from val without coupling
|
| 82 |
+
the change to train.
|
| 83 |
+
"""
|
| 84 |
+
import tensorflow as tf
|
| 85 |
+
|
| 86 |
+
load_data = _make_load_data_fn(tokenizer)
|
| 87 |
+
return (
|
| 88 |
+
tf.data.Dataset.from_tensor_slices((list(image_paths), list(captions)))
|
| 89 |
+
.map(load_data, num_parallel_calls=tf.data.AUTOTUNE)
|
| 90 |
+
.shuffle(buffer_size) # Notebook cell 14 shuffles val too — preserved.
|
| 91 |
+
.batch(batch_size)
|
| 92 |
+
)
|
src/captioning/data/splits.py
ADDED
|
@@ -0,0 +1,86 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Image-level train/val splitting.
|
| 2 |
+
|
| 3 |
+
Mirrors notebook cell 11. The split is intentionally at the *image* level
|
| 4 |
+
(not the caption level): every image owns ~5 captions in COCO, and putting
|
| 5 |
+
some of an image's captions in train and others in val would be data leakage.
|
| 6 |
+
|
| 7 |
+
The notebook does this correctly via the ``img_to_cap_vector`` defaultdict
|
| 8 |
+
loop; we preserve that exact algorithm but inject the seed so the split is
|
| 9 |
+
reproducible across runs.
|
| 10 |
+
"""
|
| 11 |
+
|
| 12 |
+
from __future__ import annotations
|
| 13 |
+
|
| 14 |
+
import collections
|
| 15 |
+
import random
|
| 16 |
+
|
| 17 |
+
import pandas as pd
|
| 18 |
+
|
| 19 |
+
from captioning.utils.logging import get_logger
|
| 20 |
+
|
| 21 |
+
log = get_logger(__name__)
|
| 22 |
+
|
| 23 |
+
|
| 24 |
+
def make_image_level_splits(
|
| 25 |
+
captions: pd.DataFrame,
|
| 26 |
+
train_fraction: float = 0.8,
|
| 27 |
+
seed: int | None = None,
|
| 28 |
+
) -> tuple[list[str], list[str], list[str], list[str]]:
|
| 29 |
+
"""Split captions into train/val while keeping all of an image's
|
| 30 |
+
captions in the same split.
|
| 31 |
+
|
| 32 |
+
Mirrors notebook cell 11 exactly when ``seed`` is the same value that was
|
| 33 |
+
fed to ``random.seed`` before the notebook ran. ``seed=None`` reproduces
|
| 34 |
+
the notebook's non-deterministic behaviour.
|
| 35 |
+
|
| 36 |
+
Args:
|
| 37 |
+
captions: DataFrame with ``image`` and ``caption`` columns
|
| 38 |
+
(preprocessed if you want preprocessed splits — the loader applies
|
| 39 |
+
``preprocess_caption`` upstream).
|
| 40 |
+
train_fraction: Fraction of *unique images* assigned to the train
|
| 41 |
+
split. The notebook uses ``int(len(img_keys) * 0.8)``, which we
|
| 42 |
+
preserve byte-for-byte (``int()`` truncates, not rounds).
|
| 43 |
+
seed: If provided, used to seed Python's ``random`` for the shuffle.
|
| 44 |
+
|
| 45 |
+
Returns:
|
| 46 |
+
Tuple ``(train_imgs, train_captions, val_imgs, val_captions)`` where
|
| 47 |
+
each list has one entry per (image, caption) pair, expanded so an
|
| 48 |
+
image with N captions appears N times.
|
| 49 |
+
"""
|
| 50 |
+
img_to_cap = collections.defaultdict(list)
|
| 51 |
+
for img, cap in zip(captions["image"], captions["caption"], strict=True):
|
| 52 |
+
img_to_cap[img].append(cap)
|
| 53 |
+
|
| 54 |
+
img_keys = list(img_to_cap.keys())
|
| 55 |
+
|
| 56 |
+
if seed is not None:
|
| 57 |
+
rng = random.Random(seed) # — seeded RNG is reproducible by design
|
| 58 |
+
rng.shuffle(img_keys)
|
| 59 |
+
else:
|
| 60 |
+
random.shuffle(img_keys)
|
| 61 |
+
|
| 62 |
+
slice_index = int(len(img_keys) * train_fraction)
|
| 63 |
+
train_keys, val_keys = img_keys[:slice_index], img_keys[slice_index:]
|
| 64 |
+
|
| 65 |
+
train_imgs: list[str] = []
|
| 66 |
+
train_captions: list[str] = []
|
| 67 |
+
for k in train_keys:
|
| 68 |
+
n = len(img_to_cap[k])
|
| 69 |
+
train_imgs.extend([k] * n)
|
| 70 |
+
train_captions.extend(img_to_cap[k])
|
| 71 |
+
|
| 72 |
+
val_imgs: list[str] = []
|
| 73 |
+
val_captions: list[str] = []
|
| 74 |
+
for k in val_keys:
|
| 75 |
+
n = len(img_to_cap[k])
|
| 76 |
+
val_imgs.extend([k] * n)
|
| 77 |
+
val_captions.extend(img_to_cap[k])
|
| 78 |
+
|
| 79 |
+
log.info(
|
| 80 |
+
"splits_made",
|
| 81 |
+
train_images=len(train_keys),
|
| 82 |
+
val_images=len(val_keys),
|
| 83 |
+
train_captions=len(train_captions),
|
| 84 |
+
val_captions=len(val_captions),
|
| 85 |
+
)
|
| 86 |
+
return train_imgs, train_captions, val_imgs, val_captions
|
src/captioning/evaluation/__init__.py
ADDED
|
@@ -0,0 +1,90 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Evaluation — caption-quality metrics + per-sample diagnostics.
|
| 2 |
+
|
| 3 |
+
Available metrics (all corpus-level, 0-100 scale where applicable):
|
| 4 |
+
* BLEU-1..4 — :mod:`bleu`
|
| 5 |
+
* ROUGE-L — :mod:`rouge`
|
| 6 |
+
* METEOR — :mod:`meteor` (requires a JRE on PATH)
|
| 7 |
+
* CIDEr — :mod:`cider` (requires >= 2 examples)
|
| 8 |
+
|
| 9 |
+
:func:`compute_all_metrics` in :mod:`runner` is the single entry point used
|
| 10 |
+
by the CLI and by future Phase 3 benchmark comparisons; per-sample
|
| 11 |
+
diagnostics live in :mod:`inspection`. Phase 3 latency timing lives in
|
| 12 |
+
:mod:`latency`, and the dashboard data export in :mod:`dashboard`.
|
| 13 |
+
"""
|
| 14 |
+
|
| 15 |
+
from captioning.evaluation.benchmark import RunMeta, write_run_artifacts
|
| 16 |
+
from captioning.evaluation.bleu import (
|
| 17 |
+
BleuBreakdown,
|
| 18 |
+
corpus_bleu_breakdown,
|
| 19 |
+
corpus_bleu_score,
|
| 20 |
+
)
|
| 21 |
+
from captioning.evaluation.cider import MIN_SAMPLES_FOR_CIDER, corpus_cider_score
|
| 22 |
+
from captioning.evaluation.comparison import (
|
| 23 |
+
ComparisonError,
|
| 24 |
+
RunRecord,
|
| 25 |
+
build_summary,
|
| 26 |
+
load_run,
|
| 27 |
+
render_markdown,
|
| 28 |
+
)
|
| 29 |
+
from captioning.evaluation.dashboard import (
|
| 30 |
+
DashboardExportError,
|
| 31 |
+
build_dashboard_data,
|
| 32 |
+
render_dashboard_json,
|
| 33 |
+
)
|
| 34 |
+
from captioning.evaluation.inspection import (
|
| 35 |
+
SampleDiagnostics,
|
| 36 |
+
diagnose_many,
|
| 37 |
+
diagnose_sample,
|
| 38 |
+
format_diagnostic_row,
|
| 39 |
+
write_diagnostics_jsonl,
|
| 40 |
+
)
|
| 41 |
+
from captioning.evaluation.latency import (
|
| 42 |
+
BatchLatency,
|
| 43 |
+
LatencyBenchmarkError,
|
| 44 |
+
LatencySettings,
|
| 45 |
+
measure_latency,
|
| 46 |
+
runtime_info,
|
| 47 |
+
summarize,
|
| 48 |
+
time_load,
|
| 49 |
+
)
|
| 50 |
+
from captioning.evaluation.meteor import corpus_meteor_score
|
| 51 |
+
from captioning.evaluation.rouge import corpus_rouge_l_score
|
| 52 |
+
from captioning.evaluation.runner import MetricsReport, compute_all_metrics
|
| 53 |
+
from captioning.evaluation.slice import EvalSlice, load_eval_slice, slice_fingerprint
|
| 54 |
+
|
| 55 |
+
__all__ = [
|
| 56 |
+
"MIN_SAMPLES_FOR_CIDER",
|
| 57 |
+
"BatchLatency",
|
| 58 |
+
"BleuBreakdown",
|
| 59 |
+
"ComparisonError",
|
| 60 |
+
"DashboardExportError",
|
| 61 |
+
"EvalSlice",
|
| 62 |
+
"LatencyBenchmarkError",
|
| 63 |
+
"LatencySettings",
|
| 64 |
+
"MetricsReport",
|
| 65 |
+
"RunMeta",
|
| 66 |
+
"RunRecord",
|
| 67 |
+
"SampleDiagnostics",
|
| 68 |
+
"build_dashboard_data",
|
| 69 |
+
"build_summary",
|
| 70 |
+
"compute_all_metrics",
|
| 71 |
+
"corpus_bleu_breakdown",
|
| 72 |
+
"corpus_bleu_score",
|
| 73 |
+
"corpus_cider_score",
|
| 74 |
+
"corpus_meteor_score",
|
| 75 |
+
"corpus_rouge_l_score",
|
| 76 |
+
"diagnose_many",
|
| 77 |
+
"diagnose_sample",
|
| 78 |
+
"format_diagnostic_row",
|
| 79 |
+
"load_eval_slice",
|
| 80 |
+
"load_run",
|
| 81 |
+
"measure_latency",
|
| 82 |
+
"render_dashboard_json",
|
| 83 |
+
"render_markdown",
|
| 84 |
+
"runtime_info",
|
| 85 |
+
"slice_fingerprint",
|
| 86 |
+
"summarize",
|
| 87 |
+
"time_load",
|
| 88 |
+
"write_diagnostics_jsonl",
|
| 89 |
+
"write_run_artifacts",
|
| 90 |
+
]
|
src/captioning/evaluation/benchmark.py
ADDED
|
@@ -0,0 +1,128 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Benchmark-ready run artefacts.
|
| 2 |
+
|
| 3 |
+
Every evaluation pass writes a consistent set of files under
|
| 4 |
+
``<run_root>/<run_id>/`` so Phase 3 cross-model comparisons can join them
|
| 5 |
+
without bespoke parsing per model:
|
| 6 |
+
|
| 7 |
+
metrics.json — :class:`MetricsReport` dumped via dataclass-asdict
|
| 8 |
+
predictions.jsonl — one row per (image, prediction, references)
|
| 9 |
+
diagnostics.jsonl — one :class:`SampleDiagnostics` per row
|
| 10 |
+
run_meta.json — model id, decode strategy, n_samples, timestamp
|
| 11 |
+
report.md — Markdown summary humans actually read
|
| 12 |
+
|
| 13 |
+
A "run" is one (model, decode_strategy, dataset_slice) tuple. ``run_id`` is
|
| 14 |
+
a free-form string — the CLI defaults to a timestamp; comparison code groups
|
| 15 |
+
by ``model_id`` to plot bars across models.
|
| 16 |
+
"""
|
| 17 |
+
|
| 18 |
+
from __future__ import annotations
|
| 19 |
+
|
| 20 |
+
import json
|
| 21 |
+
from dataclasses import dataclass, field
|
| 22 |
+
from datetime import datetime, timezone
|
| 23 |
+
from pathlib import Path
|
| 24 |
+
|
| 25 |
+
from captioning.evaluation.inspection import SampleDiagnostics, write_diagnostics_jsonl
|
| 26 |
+
from captioning.evaluation.runner import MetricsReport
|
| 27 |
+
|
| 28 |
+
|
| 29 |
+
@dataclass(frozen=True)
|
| 30 |
+
class RunMeta:
|
| 31 |
+
"""Per-evaluation-run metadata persisted next to metrics."""
|
| 32 |
+
|
| 33 |
+
model_id: str
|
| 34 |
+
decode_strategy: str
|
| 35 |
+
weights_path: str
|
| 36 |
+
tokenizer_dir: str
|
| 37 |
+
n_samples: int
|
| 38 |
+
max_length: int
|
| 39 |
+
beam_width: int | None = None
|
| 40 |
+
length_penalty: float | None = None
|
| 41 |
+
repetition_penalty: float | None = None
|
| 42 |
+
timestamp_utc: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat())
|
| 43 |
+
|
| 44 |
+
def to_dict(self) -> dict[str, object]:
|
| 45 |
+
return {
|
| 46 |
+
"model_id": self.model_id,
|
| 47 |
+
"decode_strategy": self.decode_strategy,
|
| 48 |
+
"weights_path": self.weights_path,
|
| 49 |
+
"tokenizer_dir": self.tokenizer_dir,
|
| 50 |
+
"n_samples": self.n_samples,
|
| 51 |
+
"max_length": self.max_length,
|
| 52 |
+
"beam_width": self.beam_width,
|
| 53 |
+
"length_penalty": self.length_penalty,
|
| 54 |
+
"repetition_penalty": self.repetition_penalty,
|
| 55 |
+
"timestamp_utc": self.timestamp_utc,
|
| 56 |
+
}
|
| 57 |
+
|
| 58 |
+
|
| 59 |
+
def write_run_artifacts(
|
| 60 |
+
run_dir: str | Path,
|
| 61 |
+
*,
|
| 62 |
+
metrics: MetricsReport,
|
| 63 |
+
meta: RunMeta,
|
| 64 |
+
images: list[str],
|
| 65 |
+
predictions: list[str],
|
| 66 |
+
references: list[list[str]],
|
| 67 |
+
diagnostics: list[SampleDiagnostics],
|
| 68 |
+
) -> Path:
|
| 69 |
+
"""Write every benchmark artefact to ``run_dir`` and return the directory.
|
| 70 |
+
|
| 71 |
+
Idempotent over a clean ``run_dir``; overwrites existing files inside.
|
| 72 |
+
"""
|
| 73 |
+
out = Path(run_dir)
|
| 74 |
+
out.mkdir(parents=True, exist_ok=True)
|
| 75 |
+
|
| 76 |
+
(out / "metrics.json").write_text(json.dumps(metrics.to_dict(), indent=2), encoding="utf-8")
|
| 77 |
+
(out / "run_meta.json").write_text(json.dumps(meta.to_dict(), indent=2), encoding="utf-8")
|
| 78 |
+
|
| 79 |
+
with (out / "predictions.jsonl").open("w", encoding="utf-8") as f:
|
| 80 |
+
for img, pred, refs in zip(images, predictions, references, strict=True):
|
| 81 |
+
row = {"image": img, "prediction": pred, "references": list(refs)}
|
| 82 |
+
f.write(json.dumps(row, ensure_ascii=False) + "\n")
|
| 83 |
+
|
| 84 |
+
write_diagnostics_jsonl(diagnostics, out / "diagnostics.jsonl")
|
| 85 |
+
(out / "report.md").write_text(_render_report_markdown(meta, metrics), encoding="utf-8")
|
| 86 |
+
return out
|
| 87 |
+
|
| 88 |
+
|
| 89 |
+
def _render_report_markdown(meta: RunMeta, m: MetricsReport) -> str:
|
| 90 |
+
"""Render the human-facing Markdown summary of a single run."""
|
| 91 |
+
|
| 92 |
+
def fmt(v: float | None) -> str:
|
| 93 |
+
return "n/a" if v is None else f"{v:.2f}"
|
| 94 |
+
|
| 95 |
+
lines = [
|
| 96 |
+
f"# Evaluation run — {meta.model_id}",
|
| 97 |
+
"",
|
| 98 |
+
f"- Decode strategy: `{meta.decode_strategy}`",
|
| 99 |
+
f"- Weights: `{meta.weights_path}`",
|
| 100 |
+
f"- Tokenizer dir: `{meta.tokenizer_dir}`",
|
| 101 |
+
f"- Samples: **{meta.n_samples}**",
|
| 102 |
+
f"- Timestamp (UTC): {meta.timestamp_utc}",
|
| 103 |
+
]
|
| 104 |
+
if meta.beam_width is not None:
|
| 105 |
+
lines.append(f"- Beam width: {meta.beam_width}")
|
| 106 |
+
if meta.length_penalty is not None:
|
| 107 |
+
lines.append(f"- Length penalty: {meta.length_penalty}")
|
| 108 |
+
if meta.repetition_penalty is not None:
|
| 109 |
+
lines.append(f"- Repetition penalty: {meta.repetition_penalty}")
|
| 110 |
+
lines += [
|
| 111 |
+
"",
|
| 112 |
+
"## Metrics",
|
| 113 |
+
"",
|
| 114 |
+
"| Metric | Value |",
|
| 115 |
+
"|---|---|",
|
| 116 |
+
f"| BLEU-1 | {fmt(m.bleu1)} |",
|
| 117 |
+
f"| BLEU-2 | {fmt(m.bleu2)} |",
|
| 118 |
+
f"| BLEU-3 | {fmt(m.bleu3)} |",
|
| 119 |
+
f"| BLEU-4 | {fmt(m.bleu4)} |",
|
| 120 |
+
f"| ROUGE-L | {fmt(m.rouge_l)} |",
|
| 121 |
+
f"| METEOR | {fmt(m.meteor)} |",
|
| 122 |
+
f"| CIDEr | {fmt(m.cider)} |",
|
| 123 |
+
]
|
| 124 |
+
if m.errors:
|
| 125 |
+
lines += ["", "## Skipped or failed metrics", ""]
|
| 126 |
+
for name, err in m.errors.items():
|
| 127 |
+
lines.append(f"- `{name}`: {err}")
|
| 128 |
+
return "\n".join(lines) + "\n"
|