fix: cap token×beam product so generation fits its GPU reservation

_estimate_duration clamps the ZeroGPU reservation at 120s, but generation was
uncapped, so an in-bounds worst case (1024 tokens × 8 beams) could run past the
reservation and be killed mid-decode. Cap max_new_tokens × num_beams at 720 in
_normalize_params (the single funnel translate() and _estimate_duration share),
trimming the token count — not beam width, which drives quality — so the
estimate (30 + product//8) never exceeds 120s and generation stays within budget.

Tests: add test_normalize_params_caps_token_beam_product; update the clamp test
for the cap. 69 fast pass. Docs synced to 79/69/10.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

CLAUDE.md

@@ -25,8 +25,8 @@ uv run ruff format .
 uv run ty check
 
 # Test
-uv run pytest                     # all 78 tests (slow require CUDA + model download)
-uv run pytest -m "not slow"       # 68 fast tests only
+uv run pytest                     # all 79 tests (slow require CUDA + model download)
+uv run pytest -m "not slow"       # 69 fast tests only
 uv run pytest -m slow             # 10 model tests only (CUDA only)
 
 # Generate language mapping (dev only)
@@ -35,13 +35,13 @@ uv run scripts/generate_langmap.py <path-to-paper.pdf>
 
 ## Architecture
 
-**`app.py`** — Single-file application with a Google Translate-style layout: top row has two symmetric, filterable, region-sorted language dropdowns (source defaults to "English (en)", target defaults to "French (fr)") with a swap button ("⇄") between them; below that, input textbox (autofocused) and output textbox with copy button side by side. The Translate button spans full width below both textboxes (shows "Translating..." during processing). Ctrl+Enter submits from the input. The model auto-detects source language; the source dropdown is for user reference and the swap button only, which an `info=` caption discloses. Each control carries an `info=` caption (caption text, not HTML/Markdown blocks): the target dropdown a quality-varies caveat, the input the Ctrl+Enter hint, the output model/arXiv/license provenance. Uses `@lru_cache` for lazy loading of the `google/madlad400-3b-mt` tokenizer and model. On ZeroGPU (`SPACES_ZERO_GPU=1`), `_maybe_eager_load()` places the model at module scope so the `spaces` hijack can pack weights and stream them into workers for fast cold starts; off-ZeroGPU (local, tests, cpu-basic) it stays lazy, so importing the app never downloads the model. Uses `bfloat16` on CUDA (T5/MADLAD is numerically unstable in `float16` — fp16's narrow range overflows to inf/NaN; bf16 is the format T5 was trained in), `float32` on CPU. MPS is not supported (produces garbage output with T5 models). Translation prepends a target language token with a space to the input text (e.g., `<2fr> Hello`) before tokenization and generation; whitespace-only or `None` input short-circuits to an empty string before the model loads. The generation params are normalized in `translate()` via `_normalize_params` (`None`/`NaN` → default, then clamped to range) so the cast-less public path and the ZeroGPU duration callable can't crash on a cleared `gr.Number` field. Decoding is greedy by default (deterministic); a non-default `temperature` (tolerance-compared to absorb float spinner drift) enables sampling, and `num_beams > 1` uses beam search. A collapsed "Advanced" accordion exposes `max_new_tokens`/`num_beams`/`temperature` as `gr.Number` controls (no sliders; defaults mirror `translate()`, so the default surface stays greedy). Right-to-left target scripts (an explicit `RTL_CODES` token set — `region` is not a usable proxy) flip the output box to RTL via the Translate-button and swap paths; Ctrl+Enter/`/translate` return a bare string and stay LTR. The `@spaces.GPU` decorator allocates GPU on HF Spaces infrastructure; its `duration` is a callable (`_estimate_duration`) that scales the GPU reservation with `max_new_tokens × num_beams` (capped at 120s). Both translate handlers (the private Translate-button click and the public submit) carry the advanced params, so Ctrl+Enter and the `/translate` API honor the accordion; the params keep defaults, so existing two-arg callers still work. The submit handler exposes a stable `/translate` API endpoint (returns a bare string); the swap and Translate-button handlers are `api_visibility="private"`, and both generation handlers use `show_progress="minimal"`. Only `/translate` is public.
+**`app.py`** — Single-file application with a Google Translate-style layout: top row has two symmetric, filterable, region-sorted language dropdowns (source defaults to "English (en)", target defaults to "French (fr)") with a swap button ("⇄") between them; below that, input textbox (autofocused) and output textbox with copy button side by side. The Translate button spans full width below both textboxes (shows "Translating..." during processing). Ctrl+Enter submits from the input. The model auto-detects source language; the source dropdown is for user reference and the swap button only, which an `info=` caption discloses. Each control carries an `info=` caption (caption text, not HTML/Markdown blocks): the target dropdown a quality-varies caveat, the input the Ctrl+Enter hint, the output model/arXiv/license provenance. Uses `@lru_cache` for lazy loading of the `google/madlad400-3b-mt` tokenizer and model. On ZeroGPU (`SPACES_ZERO_GPU=1`), `_maybe_eager_load()` places the model at module scope so the `spaces` hijack can pack weights and stream them into workers for fast cold starts; off-ZeroGPU (local, tests, cpu-basic) it stays lazy, so importing the app never downloads the model. Uses `bfloat16` on CUDA (T5/MADLAD is numerically unstable in `float16` — fp16's narrow range overflows to inf/NaN; bf16 is the format T5 was trained in), `float32` on CPU. MPS is not supported (produces garbage output with T5 models). Translation prepends a target language token with a space to the input text (e.g., `<2fr> Hello`) before tokenization and generation; whitespace-only or `None` input short-circuits to an empty string before the model loads. The generation params are normalized in `translate()` via `_normalize_params` (`None`/`NaN` → default, then clamped to range) so the cast-less public path and the ZeroGPU duration callable can't crash on a cleared `gr.Number` field; it also caps the `max_new_tokens × num_beams` product (`_MAX_TOKEN_BEAM_PRODUCT = 720`, trimming the token count) so a request can't outlive the GPU time `_estimate_duration` reserves. Decoding is greedy by default (deterministic); a non-default `temperature` (tolerance-compared to absorb float spinner drift) enables sampling, and `num_beams > 1` uses beam search. A collapsed "Advanced" accordion exposes `max_new_tokens`/`num_beams`/`temperature` as `gr.Number` controls (no sliders; defaults mirror `translate()`, so the default surface stays greedy). Right-to-left target scripts (an explicit `RTL_CODES` token set — `region` is not a usable proxy) flip the output box to RTL via the Translate-button and swap paths; Ctrl+Enter/`/translate` return a bare string and stay LTR. The `@spaces.GPU` decorator allocates GPU on HF Spaces infrastructure; its `duration` is a callable (`_estimate_duration`) that scales the GPU reservation with `max_new_tokens × num_beams` (capped at 120s). Both translate handlers (the private Translate-button click and the public submit) carry the advanced params, so Ctrl+Enter and the `/translate` API honor the accordion; the params keep defaults, so existing two-arg callers still work. The submit handler exposes a stable `/translate` API endpoint (returns a bare string); the swap and Translate-button handlers are `api_visibility="private"`, and both generation handlers use `show_progress="minimal"`. Only `/translate` is public.
 
 **`langmap/`** — Package with `langid_mapping.py`, mapping 418 language tokens to `{"name": ..., "region": ...}` dicts. Auto-generated by `scripts/generate_langmap.py` from Table 9 (Section A.1) of the MADLAD-400 paper. Available languages at runtime are the intersection of this mapping and the model's vocabulary.
 
 **`scripts/`** — `generate_langmap.py` parses the MADLAD-400 paper PDF (Table 9, pages 16-22) using pdfplumber and generates the static language mapping with region assignments. Dev-only tool; requires `requirements-dev.txt` dependencies.
 
-**`tests/`** — 78 tests (68 fast, 10 slow). `test_langmap.py` has 10 fast tests for mapping validation (dict shape, regions, spot-checks). `test_app.py` has 58 fast tests (signatures, device fallback, bfloat16/float32 dtype selection, ZeroGPU eager-load gating, GPU duration estimator and its signature-mirror contract + `None`-safety, greedy-by-default decoding, param forwarding into `generate`, `_normalize_params` None/NaN/clamp coercion, empty/`None`-input short-circuit, RTL output direction on the button and swap paths, `RTL_CODES` ⊆ langmap invariant, `requirements.txt` excludes platform packages, UI layout with symmetric dropdowns, swap button, textbox config including toolbar buttons and input autofocus, `info=` captions on dropdowns and textboxes spot-checked by content, the Advanced accordion's `gr.Number` controls and their bounds, advanced params reaching the public endpoint by `api_visibility` with the `/translate` input order pinned by label, `show_progress="minimal"` on generation handlers, handler wiring, stable `translate` API endpoint carrying the advanced params with UI-only handlers kept private, no HTML elements, no sliders, locale codes, no title) and 10 slow tests (translation with various parameters, language mapping). Slow tests require CUDA and model download; auto-skipped without CUDA.
+**`tests/`** — 79 tests (69 fast, 10 slow). `test_langmap.py` has 10 fast tests for mapping validation (dict shape, regions, spot-checks). `test_app.py` has 59 fast tests (signatures, device fallback, bfloat16/float32 dtype selection, ZeroGPU eager-load gating, GPU duration estimator and its signature-mirror contract + `None`-safety, greedy-by-default decoding, param forwarding into `generate`, `_normalize_params` None/NaN/clamp coercion and token×beam product cap, empty/`None`-input short-circuit, RTL output direction on the button and swap paths, `RTL_CODES` ⊆ langmap invariant, `requirements.txt` excludes platform packages, UI layout with symmetric dropdowns, swap button, textbox config including toolbar buttons and input autofocus, `info=` captions on dropdowns and textboxes spot-checked by content, the Advanced accordion's `gr.Number` controls and their bounds, advanced params reaching the public endpoint by `api_visibility` with the `/translate` input order pinned by label, `show_progress="minimal"` on generation handlers, handler wiring, stable `translate` API endpoint carrying the advanced params with UI-only handlers kept private, no HTML elements, no sliders, locale codes, no title) and 10 slow tests (translation with various parameters, language mapping). Slow tests require CUDA and model download; auto-skipped without CUDA.
 
 ## Tooling
 

README.md

@@ -39,6 +39,6 @@ The Gradio interface launches at `http://localhost:7860`.
 uv run ruff check .             # lint
 uv run ruff format .            # format
 uv run ty check                 # type check
-uv run pytest -m "not slow"     # 68 fast tests
-uv run pytest                   # all 78 tests (slow require CUDA + model download)
+uv run pytest -m "not slow"     # 69 fast tests
+uv run pytest                   # all 79 tests (slow require CUDA + model download)
 ```

app.py

@@ -106,6 +106,13 @@ def _maybe_eager_load() -> None:
         _load_model()
 
 
+# Cap the token×beam product so even the largest request fits the GPU time _estimate_duration
+# reserves. That estimate is 30 + product//8 (clamped at 120s), so a product <= 720 keeps it at
+# <= 120s, meaning generation can't outlive its reservation and be killed mid-decode. The token
+# count is trimmed to honour this (not the beam width, which drives translation quality).
+_MAX_TOKEN_BEAM_PRODUCT = 720
+
+
 def _normalize_params(
     max_new_tokens: float | None, num_beams: float | None, temperature: float | None
 ) -> tuple[int, int, float]:
@@ -113,16 +120,18 @@ def _normalize_params(
     ``None`` (Gradio skips its bounds check for ``None``) and the public ``/translate`` path
     passes values uncast, so this is the single funnel every caller — button, submit, API,
     direct, and the ZeroGPU duration callable — goes through. ``None``/``NaN`` fall back to the
-    defaults; values are clamped to the ranges the Advanced ``gr.Number`` controls advertise."""
+    defaults; values are clamped to the ranges the Advanced ``gr.Number`` controls advertise; and
+    the token×beam product is capped (``_MAX_TOKEN_BEAM_PRODUCT``) so generation stays within the
+    GPU time it reserved."""
 
     def _num(value: float | None, default: float) -> float:
         return default if value is None or math.isnan(value) else value
 
-    return (
-        int(max(1, min(1024, _num(max_new_tokens, 512)))),
-        int(max(1, min(8, _num(num_beams, 1)))),
-        float(max(0.1, min(2.0, _num(temperature, 1.0)))),
-    )
+    beams = int(max(1, min(8, _num(num_beams, 1))))
+    tokens = int(max(1, min(1024, _num(max_new_tokens, 512))))
+    tokens = min(tokens, _MAX_TOKEN_BEAM_PRODUCT // beams)
+    temperature = float(max(0.1, min(2.0, _num(temperature, 1.0))))
+    return tokens, beams, temperature
 
 
 def _estimate_duration(

tests/test_app.py

@@ -206,13 +206,27 @@ def test_normalize_params_clamps_and_defaults():
     assert app._normalize_params(None, None, None) == (512, 1, 1.0)
     nan = float("nan")
     assert app._normalize_params(nan, nan, nan) == (512, 1, 1.0)
-    assert app._normalize_params(99999, 99, 9.0) == (1024, 8, 2.0)  # clamp high
     assert app._normalize_params(0, 0, 0.0) == (1, 1, 0.1)  # clamp low
+    assert app._normalize_params(99999, 1, 9.0) == (720, 1, 2.0)  # tokens to product budget; temp clamped
+    assert app._normalize_params(10, 99, 1.0) == (10, 8, 1.0)  # beams clamp high (product 80 within budget)
     mnt, beams, temp = app._normalize_params(10.0, 4.0, 0.5)
     assert (mnt, beams, temp) == (10, 4, 0.5)
     assert type(mnt) is int and type(beams) is int and type(temp) is float
 
 
+def test_normalize_params_caps_token_beam_product():
+    """High token×beam requests get the token count trimmed so the product stays within the GPU
+    reservation _estimate_duration grants (the worst case must not outlive its 120s budget)."""
+    import app
+
+    mnt, beams, _ = app._normalize_params(1024, 8, 1.0)
+    assert beams == 8 and mnt * beams <= app._MAX_TOKEN_BEAM_PRODUCT
+    # the duration estimate, built on the same normalization, never exceeds its 120s cap
+    assert app._estimate_duration("hi", "French (fr)", 1024, 8, 1.0) <= 120
+    # comfortably-small products are left untouched
+    assert app._normalize_params(200, 2, 1.0) == (200, 2, 1.0)
+
+
 def test_translate_normalizes_invalid_params():
     """A cleared gr.Number arrives as None (and temperature can be NaN) on the public path;
     translate() must coerce to defaults instead of crashing or corrupting sampling."""