subtitle-linebreak-bilingual

<h1 align="center">subtitle-linebreak-bilingual</h1> <p align="center"> <b><a href="https://huggingface.co/suko/subtitle-linebreak-bilingual/blob/main/README.zh-Hant.md">繁體中文</a></b> </p> <p align="center"> 42MB FP32 ONNX / 10.8MB INT8 PyTorch · 100% local inference (browser/CPU) · Chinese + English (incl. code-switched) · no LLM, no API calls </p>

A bilingual (Traditional Chinese + English) classifier that predicts, for a full sentence of transcribed speech, where a subtitle line break should go — the same decision a human caption editor makes when splitting a long sentence into display lines. Built as a local, zero-API-cost alternative to sending transcript text to an LLM for this decision.

One model file covers both languages, including code-switched Chinese/English sentences, which is the common case in Taiwanese YouTube content.

Architecture

• Base encoder: voidful/albert_chinese_base — ALBERT, 12 layers, hidden size 768, 10.55M parameters thanks to cross-layer weight sharing.
• Case-insensitive tokenizer (do_lower_case=True): capitalized English tokenizes into ordinary wordpieces, not [UNK] — the property that makes a single bilingual model viable at this size.
• Head: token representation + one scalar feature (remaining line-length budget at that position), through a small classification head.

Candidate break units

The sentence is split into break-candidate units before classification:

• each CJK character is one unit;
• each maximal run of non-CJK, non-whitespace characters is one unit (an English word, a number, GTA6).

The joining rule used when rendering predictions back into text matches isCjkAtom in the product's own subtitle-split.ts, so training-time and inference-time segmentation agree with what ships.

Hard-constraint fallback (required for any integration)

A soft per-position classifier cannot guarantee it will never exceed a hard line-length cap, no matter how it is trained — tried adding an "urgency" penalty to the loss near the cap, which made precision worse without fully eliminating violations. The reliable fix: a deterministic decode-time fallback on top of the soft model — force a break once the cap would otherwise be exceeded, regardless of confidence. Keep this fallback in any integration.

A second artifact, orphaned punctuation lines, is suppressed in the weights. Greedy decode can fire two breaks close enough together that a punctuation-only candidate unit (a bare comma sitting between two Chinese clauses — its own one-character unit under this model's unit-splitting rule) ends up alone on its own line. Earlier releases of this model required integrators to filter this downstream. They no longer do: this checkpoint is trained with a joint penalty that suppresses the artifact directly (see Suppressing orphaned punctuation lines). infer_example.py still ships a decode-time filter for it as belt-and-suspenders, but it is now optional; the over-length fallback above is not.

Training data

Eight real, human-authored YouTube caption sources. No synthetic labels, no auto-generated captions.

• English: TED, Kurzgesagt, Molly Burke, Rikki Poynter, Philip DeFranco
• Traditional Chinese: 志祺七七, 這群人 TGOP, 小宁子 XNZ

28,639 training sentences / 3,184 validation sentences.

Contamination filter

Scraped "manual" captions aren't automatically real line-wrapped subtitles — some sources just chunk narration into cues without respecting a display-line cap, which makes those line breaks mislabeled supervision. The filter checks both constraints independently:

• MAX_LINE_VISUAL = 42 — visual width, CJK character counts as 2 columns, everything else as 1;
• DEFAULT_MAX_UNITS = 14 — word/character-segment count cap, ported verbatim from the product's textUnits().

If any part of a merged training sentence violates either cap, the whole sentence is dropped.

Performance

6 epochs. Evaluated with greedy left-to-right decoding — the model's own predictions feed the running line-width feature, not ground truth, matching real inference and giving a harder, more honest number than teacher-forced evaluation.

int8 dynamic quantization costs about 0.005 F1 across all three splits.

Artifact rates

Decoded over the full validation set (17,091 rendered lines), int8, threshold 0.5:

The three remaining punctuation-only lines are all English: two speaker-change dashes (-) and one ..., shapes that are defensible in English caption convention. Chinese, the language this artifact was reported on, is clean.

The cost is 0.9 points of Chinese F1 (−1.2% relative), and it is worth being precise about its shape: precision falls (0.6937 → 0.6825 fp32) while recall rises (0.8926 → 0.8971). The model cuts slightly more freely. Whether that trade is right for you depends on whether a missed cut has a downstream correction mechanism in your pipeline; in Cap's it does (a separate pause signal), and a wrong cut does not.

Two honesty notes on these numbers. The F1 delta is a single seed on different hardware from the previous release, so run-to-run variance is not separable from it; the artifact deltas are the same validation set under the same measurement, and are not similarly ambiguous. And on one held-out Chinese video at Cap's production threshold of 0.95, this release scores slightly worse than the previous one (F1 0.765 vs 0.777) while scoring slightly better at 0.5 (0.739 vs 0.733) — consistent with the precision/recall shift above. If you run at a high threshold, measure on your own content before switching.

Chinese essentially ties a separate, dedicated Chinese-only specialist model from this project (~30M params, single-language, not yet public: F1 = 0.7789 fp32 / 0.7752 int8) — a smaller bilingual file matching a larger single-language one. This came from an earlier finding that a smaller ALBERT base (voidful/albert_chinese_tiny, 4 layers) was capacity-limited on this task; the base sibling (same vocab/tokenizer, cross-layer weight sharing keeps the size cost modest: 15.6MB → 40.2MB fp32 for 3× the layers) closed nearly the whole gap with no new training data.

English (F1 = 0.541) is well below a dedicated English specialist (F1 = 0.65, not yet public) and a hand-written rule-based line-breaker (F1 0.79–0.84 on scripted content) — not competitive for pure English content yet. Checked whether this was actually a segmentation bug (wrong/mid-word break candidates) before assuming it's a judgment problem: it isn't. On 5 real English validation sentences, candidate units were correct whole-word boundaries with punctuation correctly attached, never split mid-word, and predictions matched human labels exactly on 4/5 — the one miss was an off-by-one word (a plausible break point, not a nonsense one). So the gap is judgment consistency (especially conversational vs. scripted register), not broken tokenization, and root cause isn't further identified. It's also not the units-cap contamination fix that drove most of the Chinese gain — English hits the 42-column visual cap before the 14-word cap, so that fix never touched English data.

Suppressing orphaned punctuation lines

The artifact — a line containing nothing but a comma — is a joint failure: it needs two breaks to fire, one before the punctuation unit and one after. Neither break is wrong on its own. That is why the obvious fix does not work.

What failed. Upweighting the BCE loss on the gaps around punctuation-only units (10×/20×/50×, with and without the second gap) cost Chinese F1 (0.6964 → 0.6495 at worst) and cost recall on breaks that should land before punctuation — opening brackets 「【《『 fell from 95.4% to 80.2% — while the artifact count itself moved only within noise. Per-gap weighting cannot express "not both". BCE sees one gap at a time; the constraint spans two.

What works. Add a term that penalizes the product of the two probabilities:

loss = BCE(logits, labels) + λ · mean over punctuation runs [a..b] of  σ(z_{a-1}) · σ(z_b^cf)

Three details carry the result:

• The product, not the sum. Penalizing the joint event leaves the model free to choose which factor to lower, and it consistently lowers the one that is not label-aligned. Penalizing either gap alone destroys the correct breaks along with the incorrect ones — that is exactly what the failed attempts did.
• z_b^cf is counterfactual. The remaining-line-budget feature for the second gap is computed as if the first break had already fired, i.e. with the line consisting of the punctuation run alone. That is the decode-time state the model will actually be in, and the state teacher forcing never shows it.
• Opening brackets are exempt. They are themselves punctuation-only units, so "break before 「" — correct, and labelled 1 — was being penalized alongside "break before a comma". This showed up as a bracket-recall cost that did not vary with λ (88.4% / 89.5% / 88.4% at λ = 0.5 / 1 / 2), which is the signature of a structural bug rather than too large a dose. Exempting opener-initial runs (22.1% of runs in training) restored bracket recall to 95.35% — above where it sat with no penalty at all — while leaving the suppression intact.

Shipped setting: λ = 1, 6 epochs. Calibrate λ on the model you are actually shipping: on a tiny encoder, λ = 2 was free on every axis; on this base encoder the same λ cost 1.5 points of Chinese F1. A small model calibrates the mechanism, not the dose.

ONNX export

An ONNX export is included (onnx/encoder.onnx + onnx/head.onnx), verified for numerical parity with the PyTorch model (max hidden-state diff 1.1e-5, logit diff 1.9e-6). fp32 only — ONNX-side int8 quantization is currently broken for this architecture. ALBERT's cross-layer weight sharing causes onnxruntime's dynamic quantizer to convert only a shared weight tensor's first consumer, leaving 90 of 97 MatMul nodes un-quantized (38.3MB → 37.0MB, no real savings); the standard fix (duplicating the shared initializer per consumer) made it dramatically worse (307.5MB) instead of better. Unresolved, flagged as open work. The .pt checkpoints below are correctly quantized via PyTorch's own dynamic quantization instead.

Input / output

Input is one plain-text sentence (merged transcript, not a single subtitle cue) — no pre-tokenization needed. Inference is two stages: (1) encode the sentence once → hidden states per token; (2) classify each candidate gap left-to-right, gathering each position's hidden state plus a runtime remaining line-length budget feature through the head. Stage 2 must run greedily, left to right — the budget resets on each break, so it can't be batched order-independently. Output: one break/no-break decision per gap (plus the raw logit, for a non-0.5 threshold).

gathered_hidden for a gap is last_hidden_state[:, tok_idx, :] (the token whose end offset matches that gap). gap_feature is (cap - current_line_width) / cap, recomputed at each greedy-decode step.

Max sequence length: 160 tokens (including [CLS]/[SEP]) — this is what the model was trained on; the base ALBERT encoder itself allows up to 512 via position embeddings, but anything the model never saw during training is out-of-distribution. If your input can exceed this, chunk it yourself before calling encoder.onnx — see infer_example.py in this repo, or Cap's own chunker at src/lib/linebreak/propose.ts for a worked example (it prefers to cut chunks at real word/pause boundaries rather than a hard token count).

Decision threshold: the performance table above uses the default 0.5. Cap's own production deployment uses 0.95 instead — measured on an unseen-channel benchmark, 0.95 gives 82% precision / 74% recall vs. 65%/84% at 0.5. Pick based on your own precision/recall tradeoff, especially whether a wrong cut has any downstream correction mechanism (Cap's does not; a missed cut does, via a separate pause-based signal). See the honesty notes under Artifact rates before running this release at a high threshold.

Files

• gap_classifier_bilingual_base_int8.pt — int8 dynamic-quantized PyTorch state dict, 10.8MB. Recommended checkpoint.
• gap_classifier_bilingual_base.pt — fp32 PyTorch state dict, 42.5MB, kept for reference and further fine-tuning.
• onnx/encoder.onnx + onnx/head.onnx — fp32 ONNX export (no working int8 ONNX yet, see above).
• infer_example.py — minimal runnable reference implementation (onnxruntime + transformers, no PyTorch needed): python infer_example.py "your sentence". Reproduces the exact preprocessing/decode logic Cap ships. Implements the required over-length fallback, plus an optional orphaned-line filter kept as belt-and-suspenders — read its docstring before building on it; it also records what did and did not work when training the suppression.
• Base encoder: voidful/albert_chinese_base

Previous release

The release before this one has no orphan suppression and slightly higher Chinese F1 (0.7792 int8). It stays reachable by commit if you want to pin or compare it:

https://huggingface.co/suko/subtitle-linebreak-bilingual/resolve/0350c6714b266028f4145959159334234e1adefb/onnx/encoder.onnx

License

Apache 2.0. This repository ships model weights only — no training data and no copyrighted caption text is redistributed.