subtitle-linebreak-bilingual

繁體中文

42MB FP32 ONNX / 10.8MB INT8 PyTorch · 100% local inference (browser/CPU) · Chinese + English (incl. code-switched) · no LLM, no API calls

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, separate guard is also required: short/orphaned lines. 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. This is not a rare edge case — it reproduces on plain, non-adversarial sentences. Cap's own pipeline filters this downstream: a line at or under half of its soft minimum unit count, with no real pause backing the cut, gets discarded. This model does not filter it itself; any integration needs its own version of this guard, in addition to the over-length one above.

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.

Split	F1 (fp32)	F1 (int8)
Overall	0.656	0.652
Chinese	0.7807	0.7792
English	0.541	0.535

int8 dynamic quantization costs essentially nothing in accuracy across all three splits.

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.

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 6.4e-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).

Graph	Input	Shape	Dtype
encoder.onnx	input_ids / attention_mask	[batch, seq_len]	int64
	→ last_hidden_state	[batch, seq_len, 768]	float32
head.onnx	gathered_hidden	[batch, 768]	float32
	gap_feature	[batch]	float32
	→ logit	[batch]	float32

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).

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, including both known raw-output caveats above (over-length and orphaned-short lines) — read its docstring before building on it.
• Base encoder: voidful/albert_chinese_base

License

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