commit ffe76e52550e49ca8ca4e51bc29c5cb65ee7c548 truex <truex@equwal.com> 2026-09-26 13:28:33 -0700 Align each page on its own with unique 8-character runs One character diff of all cues against the pages read so far took 9 s for a book of 300 000 characters. When few pages were read, it matched cues to wrong pages, and subtitles of another text matched too. The aligner now finds each page in the cue stream with anchors: the 8-character runs of the page that occur only once in the cues. A page with fewer than three anchors is not in the subtitles. A character diff compares the page with its region only. The longest rising chain of (cue character, token) pairs keeps the spans in order. The aligner keeps the result of each page, so one more page costs one page. createAligner(cues).spans(pages) is the new API. alignCuesToTokens stays as a wrapper. An unmatched cue gets an empty span. The gap fill is removed, because the UI never draws an unmatched cue. A cue stream of 64 characters or less is compared with each whole page.
README.md | 14 +- src/align.ts | 355 ++++++++++++++++++++++++++++++++------------ test/align.property.test.ts | 261 ++++++++++++++++++++++++++++---- test/align.test.ts | 213 +++++++++++++++++++------- test/book.ts | 114 ++++++++++++++ 5 files changed, 779 insertions(+), 178 deletions(-)
diff --git a/README.md b/README.md index 24ee73b..d536fd8 100644 --- a/README.md +++ b/README.md @@ -287,9 +287,17 @@ in a browser bundle. ## Scope limits -- The alignment runs over the whole book each time a page is read. It - takes about 5 s for 200k characters at 10% OCR noise, and well under - 1 s for a chapter. +- The alignment finds each page in the subtitles on its own: runs of 8 + characters that occur only once in the subtitles show where the page + is, and a character diff compares the page with that part only. So a + cue matches only a page that is read, and subtitles of another text + match nothing. A page is aligned once, when it is read. That takes a + few milliseconds; all the pages of a book of 300 000 characters take + about 1 s (Node on a desktop). +- A page with fewer than three such runs gets no cues, for example a + page with only a few words. A subtitle file of 64 characters or less + is too short for the runs: it is compared with each whole page, and + text that is not in it can match by chance. - A vertical font in a text layer (`dir: 'ttb'` in pdf.js) lays the characters down the page. This path has unit tests with mocked items, but no fixture PDF: pdf-lib does not write vertical text. diff --git a/src/align.ts b/src/align.ts index 3246457..e632e94 100644 --- a/src/align.ts +++ b/src/align.ts @@ -1,14 +1,31 @@ // Pure text alignment. No DOM, no OCR engine. // -// Input: subtitle cue texts and OCR tokens in reading order. -// Output: for each cue, the span of OCR tokens that shows that cue. +// Input: the texts of the subtitle cues, and the tokens of the pages that +// are read so far. Output: for each cue, the span of tokens that shows the +// cue. // -// Method: normalize both sides to a plain character stream. Run a -// character diff (Myers, via fast-diff) between the joined cue stream and -// the OCR stream. Equal runs of two or more characters are anchors. Each -// cue gets the span between its first and last anchor. A cue with no -// anchor gets a share of the gap between its matched neighbors. The diff -// is monotonic, so the spans are monotonic too. +// A book has some hundred thousand characters, and the reader reads its +// pages one by one. One diff of all cues against a few pages is slow, and +// it matches cues to wrong places. So each page is aligned on its own: +// +// 1. All cues, normalized and joined, make the cue stream. An index holds +// each K-gram (a run of K characters) that occurs only once in it. +// 2. Each K-gram of a page that is in the index gives an anchor: a cue +// position and a page position. The longest chain of anchors that rises +// in both stays. An anchor far from the median shift of the chain is a +// chance hit and goes too. A page with fewer than MIN_ANCHORS anchors is +// not in the subtitles. +// 3. The first and the last anchor give the region of the cue stream that +// the page shows. A character diff (Myers, via fast-diff) of the region +// against the page maps cue characters to page tokens. +// 4. The pairs (cue character, token) of all read pages must rise in both. +// The longest rising chain of pairs stays. The other pairs go, for +// example the pairs of a page that is out of order. +// 5. Each cue gets the span from its first to its last mapped token. +// +// The aligner keeps the result of steps 2 and 3 for each page, so one more +// page costs the work of one page. A short cue stream has too few K-grams +// for anchors: its region is the whole stream. import diff from 'fast-diff'; @@ -23,23 +40,40 @@ export interface OcrToken { bbox: { x0: number; y0: number; x1: number; y1: number }; } -/** Half-open character span [start, end) in the normalized OCR stream. */ -export interface CharSpan { - start: number; - end: number; - /** True when at least one anchor supports the span. */ - matched: boolean; -} - -/** Half-open token span [start, end) in the token list. */ +/** Half-open token span [start, end) in the tokens of the read pages. */ export interface TokenSpan { start: number; end: number; + /** True when a character of the cue is mapped. An unmatched cue has the span [0, 0). */ matched: boolean; } -/** Shortest equal run that counts as an anchor. */ -const MIN_ANCHOR = 2; +/** The length of the K-grams in the index. */ +const K = 8; + +/** A page with fewer anchors than this is not in the subtitles. */ +const MIN_ANCHORS = 3; + +/** + * An anchor whose shift (cue position minus page position) is further than + * this from the median shift of its page is a chance hit. A long page can + * have a longer part that is not narrated, such as a caption, so a quarter + * of the page length applies when it is more. + */ +const SHIFT_LIMIT = 50; + +/** + * A cue stream of this length or less (a few subtitle lines) is too short + * for anchors. Each page is then compared with the whole stream, so text + * that is not in the subtitles can match by chance. + */ +const SHORT_STREAM = 64; + +/** The shortest equal run of the diff that counts as a match. */ +const MIN_RUN = 2; + +/** The longest edit next to a shorter equal run that still counts as local noise. */ +const MAX_LOCAL_EDIT = 2; const KEEP = /[\p{L}\p{N}\p{M}]/u; @@ -75,29 +109,175 @@ export function buildStream(texts: readonly string[]): Stream { return { chars, tokenOf }; } +/** The tokens of one page. The aligner uses only their text. */ +export type PageTokens = readonly Pick<OcrToken, 'text'>[]; + +/** Aligns the cues of one subtitle file to the pages that are read. */ +export interface Aligner { + /** + * One span for each cue. `pages[p]` holds the tokens of page p, or + * undefined when page p is not read. The spans index the tokens of the + * read pages, joined in page order. The spans of the matched cues rise: + * a later cue starts and ends at or after an earlier cue. + */ + spans(pages: readonly (PageTokens | undefined)[]): TokenSpan[]; +} + +/** The pairs (cue character, page token) of one page, in page order. Both rise. */ +interface PageMatch { + /** Positions in the cue stream. */ + cue: Int32Array; + /** For each position in `cue`, the index of its token on the page. */ + token: Int32Array; +} + +/** A half-open range [from, to) of the cue stream. */ +interface Region { + from: number; + to: number; +} + /** - * Align normalized cue texts to a normalized OCR stream. - * Returns one CharSpan per cue. Spans are monotonic and inside [0, ocr.length]. + * Make an aligner for the cues of one subtitle file. Keep it while the + * subtitles stay loaded: it keeps the result of each page that it saw. */ -export function alignCues(cueTexts: readonly string[], ocr: string): CharSpan[] { - const n = cueTexts.length; - const offsets: number[] = []; - let joined = ''; - for (const t of cueTexts) { - offsets.push(joined.length); - joined += t; +export function createAligner(cueTexts: readonly string[]): Aligner { + const cues = cueTexts.map(normalizeText); + const stream = cues.join(''); + /** cueOf[i] is the index of the cue that holds stream[i]. */ + const cueOf = new Int32Array(stream.length); + let at = 0; + cues.forEach((cue, i) => { + cueOf.fill(i, at, at + cue.length); + at += cue.length; + }); + const grams = uniqueGrams(stream); + /** The match of each page, by the identity of its token array. */ + const matches = new WeakMap<PageTokens, PageMatch>(); + + function matchPage(tokens: PageTokens): PageMatch { + const page = buildStream(tokens.map((t) => t.text)); + const region = + stream.length <= SHORT_STREAM + ? { from: 0, to: stream.length } + : locate(page.chars, stream, grams); + return region + ? mapRegion(stream, region, page) + : { cue: new Int32Array(), token: new Int32Array() }; } - offsets.push(joined.length); - // Map each anchored cue char to its OCR position. - const mapped = new Int32Array(joined.length).fill(-1); - const ops = diff(joined, ocr); - let a = 0; + return { + spans(pages) { + // Join the pairs of the read pages in page order. A token index counts + // the tokens of all read pages before it. + const parts: { match: PageMatch; base: number }[] = []; + let base = 0; + let count = 0; + for (const tokens of pages) { + if (!tokens) continue; + let match = matches.get(tokens); + if (!match) { + match = matchPage(tokens); + matches.set(tokens, match); + } + parts.push({ match, base }); + base += tokens.length; + count += match.cue.length; + } + const cue = new Int32Array(count); + const token = new Int32Array(count); + let n = 0; + for (const part of parts) { + cue.set(part.match.cue, n); + for (let k = 0; k < part.match.token.length; k++) { + token[n + k] = part.base + part.match.token[k]!; + } + n += part.match.cue.length; + } + + // The tokens rise already. Keep the longest chain of pairs whose cue + // positions rise too. Then each cue spans its first to its last token. + const spans: TokenSpan[] = cues.map(() => ({ start: 0, end: 0, matched: false })); + for (const k of risingChain(cue)) { + const span = spans[cueOf[cue[k]!]!]!; + if (!span.matched) { + span.start = token[k]!; + span.matched = true; + } + span.end = token[k]! + 1; + } + return spans; + }, + }; +} + +/** A hash of the K characters of `text` at `at`. Two different K-grams can have the same hash. */ +function gramHash(text: string, at: number): number { + let h = 0; + for (let i = at; i < at + K; i++) h = (Math.imul(h, 31) + text.charCodeAt(i)) | 0; + return h; +} + +/** + * For each K-gram hash of the stream: the position of the K-gram when the + * hash occurs once, else -1. + */ +function uniqueGrams(stream: string): Map<number, number> { + const grams = new Map<number, number>(); + for (let i = 0; i + K <= stream.length; i++) { + const h = gramHash(stream, i); + grams.set(h, grams.has(h) ? -1 : i); + } + return grams; +} + +/** + * The region of the cue stream that a page shows, found with anchors. + * Null when the page has too few anchors: it is not in the subtitles. + */ +function locate(page: string, stream: string, grams: ReadonlyMap<number, number>): Region | null { + // The anchors in page order. A hash hit counts only when the text agrees. + const cuePos: number[] = []; + const pagePos: number[] = []; + for (let j = 0; j + K <= page.length; j++) { + const at = grams.get(gramHash(page, j)); + if (at !== undefined && at >= 0 && stream.startsWith(page.slice(j, j + K), at)) { + cuePos.push(at); + pagePos.push(j); + } + } + // A chance hit breaks the order, so keep the longest chain that rises in + // both. A chance hit in the chain is far from the median shift. + const chain = risingChain(cuePos); + if (chain.length < MIN_ANCHORS) return null; + const shift = (k: number): number => cuePos[k]! - pagePos[k]!; + const median = chain.map(shift).sort((x, y) => x - y)[chain.length >> 1]!; + const limit = Math.max(SHIFT_LIMIT, page.length / 4); + const kept = chain.filter((k) => Math.abs(shift(k) - median) <= limit); + if (kept.length < MIN_ANCHORS) return null; + // The page ends, projected along the shift of the first and the last anchor. + const first = kept[0]!; + const last = kept[kept.length - 1]!; + return { + from: Math.max(0, shift(first)), + to: Math.min(stream.length, shift(last) + page.length), + }; +} + +/** Map the cue characters of a region to the tokens of a page with a character diff. */ +function mapRegion(stream: string, region: Region, page: Stream): PageMatch { + const cue: number[] = []; + const token: number[] = []; + const ops = diff(stream.slice(region.from, region.to), page.chars); + let a = region.from; let b = 0; ops.forEach(([op, text], i) => { if (op === diff.EQUAL) { - if (text.length >= MIN_ANCHOR || isLocalNoise(ops, i)) { - for (let k = 0; k < text.length; k++) mapped[a + k] = b + k; + if (text.length >= MIN_RUN || isLocalNoise(ops, i)) { + for (let k = 0; k < text.length; k++) { + cue.push(a + k); + token.push(page.tokenOf[b + k]!); + } } a += text.length; b += text.length; @@ -107,30 +287,9 @@ export function alignCues(cueTexts: readonly string[], ocr: string): CharSpan[] b += text.length; } }); - - // Build matched spans. - const spans: CharSpan[] = []; - for (let i = 0; i < n; i++) { - let lo = -1; - let hi = -1; - for (let k = offsets[i]!; k < offsets[i + 1]!; k++) { - const p = mapped[k]!; - if (p < 0) continue; - if (lo < 0) lo = p; - hi = p + 1; - } - spans.push( - lo < 0 ? { start: 0, end: 0, matched: false } : { start: lo, end: hi, matched: true }, - ); - } - - fillGaps(spans, cueTexts, ocr.length); - return spans; + return { cue: Int32Array.from(cue), token: Int32Array.from(token) }; } -/** Longest edit next to a short equal run that still counts as local noise. */ -const MAX_LOCAL_EDIT = 2; - /** * True when the equal run at ops[i] sits between short edits, such as one * substituted or one dropped character. A short equal run inside a long @@ -145,54 +304,54 @@ function isLocalNoise(ops: readonly diff.Diff[], i: number): boolean { } /** - * Give each run of unmatched cues a share of the gap between its matched - * neighbors, in proportion to cue length. + * The indices of a longest chain of `values` that rises strictly, in index + * order. Patience sorting: O(n log n). */ -function fillGaps(spans: CharSpan[], cueTexts: readonly string[], ocrLength: number): void { - let i = 0; - while (i < spans.length) { - if (spans[i]!.matched) { - i++; - continue; - } - let j = i; - while (j < spans.length && !spans[j]!.matched) j++; - const gapStart = i > 0 ? spans[i - 1]!.end : 0; - const gapEnd = j < spans.length ? spans[j]!.start : ocrLength; - let total = 0; - for (let k = i; k < j; k++) total += cueTexts[k]!.length; - let pos = gapStart; - for (let k = i; k < j; k++) { - const share = total > 0 ? Math.round(((gapEnd - gapStart) * cueTexts[k]!.length) / total) : 0; - const end = k === j - 1 ? gapEnd : Math.min(gapEnd, pos + share); - spans[k] = { start: pos, end, matched: false }; - pos = end; +export function risingChain(values: ArrayLike<number>): number[] { + /** ends[k] is the index of the smallest last value of a rising chain of length k + 1. */ + const ends = new Int32Array(values.length); + /** before[i] is the index before i in the chain that ends at i, or -1. */ + const before = new Int32Array(values.length); + let length = 0; + for (let i = 0; i < values.length; i++) { + const value = values[i]!; + // Find the shortest chain whose last value is not below `value`. Most + // values extend the longest chain, so try that first. + let lo = length; + if (length === 0 || values[ends[length - 1]!]! >= value) { + lo = 0; + let hi = length; + while (lo < hi) { + const mid = (lo + hi) >> 1; + if (values[ends[mid]!]! < value) lo = mid + 1; + else hi = mid; + } } - i = j; + before[i] = lo > 0 ? ends[lo - 1]! : -1; + ends[lo] = i; + if (lo === length) length++; } -} - -/** Convert a char span to a token span with the char-to-token map. */ -export function toTokenSpan(span: CharSpan, stream: Stream): TokenSpan { - if (span.end <= span.start) { - // Empty span: point at the token at or after the position. - const t = - span.start < stream.tokenOf.length ? stream.tokenOf[span.start]! : stream.tokenOf.length; - return { start: t, end: t, matched: span.matched }; + const chain = new Array<number>(length); + for (let k = length - 1, i = length > 0 ? ends[length - 1]! : -1; k >= 0; k--, i = before[i]!) { + chain[k] = i; } - return { - start: stream.tokenOf[span.start]!, - end: stream.tokenOf[span.end - 1]! + 1, - matched: span.matched, - }; + return chain; } -/** Full pipeline: raw cue texts and OCR tokens to token spans. */ +/** + * Align cues to tokens in one call. The tokens must be in page order, the + * order that the reader keeps them in. The spans index `tokens`. + */ export function alignCuesToTokens( cueTexts: readonly string[], - tokens: readonly Pick<OcrToken, 'text'>[], + tokens: readonly Pick<OcrToken, 'text' | 'page'>[], ): TokenSpan[] { - const stream = buildStream(tokens.map((t) => t.text)); - const spans = alignCues(cueTexts.map(normalizeText), stream.chars); - return spans.map((s) => toTokenSpan(s, stream)); + const pages: Pick<OcrToken, 'text' | 'page'>[][] = []; + tokens.forEach((t, i) => { + if (i > 0 && t.page < tokens[i - 1]!.page) { + throw new Error('alignCuesToTokens: the tokens are not in page order.'); + } + (pages[t.page] ??= []).push(t); + }); + return createAligner(cueTexts).spans(pages); } diff --git a/test/align.property.test.ts b/test/align.property.test.ts index 40cc0d0..7f51019 100644 --- a/test/align.property.test.ts +++ b/test/align.property.test.ts @@ -1,12 +1,62 @@ import fc from 'fast-check'; import { describe, expect, it } from 'vitest'; -import { alignCues, normalizeText } from '../src/align'; +import { + alignCuesToTokens, + createAligner, + normalizeText, + risingChain, + type TokenSpan, +} from '../src/align'; +import { books, char, grams, pagesOfCue, readPages, word, type Book, type Token } from './book'; -// A small mixed alphabet: Latin, kana and kanji, so CJK is always covered. -const ALPHABET = [...'abcdefghijklmnop', ...'あいうえおかきくけこ', ...'吾輩猫名前春夏秋冬']; -const char = fc.constantFrom(...ALPHABET); -const word = (min: number, max: number) => - fc.array(char, { minLength: min, maxLength: max }).map((cs) => cs.join('')); +/** One token for each character of `text`, all on `page`. */ +function tokensOf(text: string, page = 0): Token[] { + return [...text].map((ch) => ({ text: ch, page })); +} + +/** Books without noise that are long enough for anchors. */ +const exactBooks = books({ cues: [15, 60], cueLength: [5, 40], pageLength: [16, 200], noise: 0 }); + +/** A book and a random choice of read pages. About one page in three is read. */ +function withReadPages(arb: fc.Arbitrary<Book>): fc.Arbitrary<[Book, boolean[]]> { + return arb.chain((book) => + fc.tuple( + fc.constant(book), + fc.array(fc.constantFrom(true, false, false), { + minLength: book.pages.length, + maxLength: book.pages.length, + }), + ), + ); +} + +/** + * For a book without noise: the index of each character in the tokens of + * the read pages, joined in page order. -1 for a character on an unread page. + */ +function tokenIndex(book: Book, read: readonly boolean[]): number[] { + let next = 0; + return book.pageOf.map((page) => (read[page] ? next++ : -1)); +} + +/** The spans of the matched cues rise and stay inside [0, total]. */ +function expectRising(spans: readonly TokenSpan[], total: number, disjoint: boolean): void { + let last: TokenSpan | null = null; + for (const span of spans) { + if (!span.matched) { + expect(span).toEqual({ start: 0, end: 0, matched: false }); + continue; + } + expect(span.start).toBeGreaterThanOrEqual(0); + expect(span.end).toBeGreaterThan(span.start); + expect(span.end).toBeLessThanOrEqual(total); + if (last) { + expect(span.start).toBeGreaterThanOrEqual(disjoint ? last.end : last.start); + expect(span.end).toBeGreaterThanOrEqual(last.end); + } + last = span; + } +} describe('normalizeText', () => { it('is idempotent', () => { @@ -19,40 +69,198 @@ describe('normalizeText', () => { }); }); -describe('alignCues properties', () => { - it('maps exact text to exact spans', () => { +describe('risingChain', () => { + /** The length of a longest strictly rising chain, by the slow O(n²) method. */ + function slowLength(values: readonly number[]): number { + const best = values.map(() => 1); + for (let i = 0; i < values.length; i++) { + for (let j = 0; j < i; j++) { + if (values[j]! < values[i]!) best[i] = Math.max(best[i]!, best[j]! + 1); + } + } + return Math.max(0, ...best); + } + + it('gives a strictly rising chain as long as the slow reference', () => { fc.assert( - fc.property(fc.array(word(1, 30), { minLength: 1, maxLength: 20 }), (cues) => { - const spans = alignCues(cues, cues.join('')); - let offset = 0; - cues.forEach((cue, i) => { - expect(spans[i]).toMatchObject({ start: offset, end: offset + cue.length }); - offset += cue.length; + fc.property(fc.array(fc.integer({ min: 0, max: 20 }), { maxLength: 40 }), (values) => { + const chain = risingChain(values); + for (let k = 1; k < chain.length; k++) { + expect(chain[k]!).toBeGreaterThan(chain[k - 1]!); + expect(values[chain[k]!]!).toBeGreaterThan(values[chain[k - 1]!]!); + } + expect(chain.length).toBe(slowLength(values)); + }), + ); + }); +}); + +describe('createAligner properties', () => { + it('maps exact text on all pages to exactly the tokens of each cue', () => { + fc.assert( + fc.property(exactBooks, (book) => { + const spans = createAligner(book.cues).spans(book.pages); + book.cues.forEach((_, i) => { + expect(spans[i]).toEqual({ + start: book.offsets[i], + end: book.offsets[i + 1], + matched: true, + }); }); }), ); }); - it('gives every cue a contiguous in-range span, monotonic across cues', () => { + it('matches cues only on read pages, and each cue there to exactly its tokens', () => { + // The regression: with a few pages read, cues matched to wrong places. + fc.assert( + fc.property(withReadPages(exactBooks), ([book, read]) => { + const spans = createAligner(book.cues).spans(readPages(book, read)); + const index = tokenIndex(book, read); + book.cues.forEach((_, i) => { + const tokens = index.slice(book.offsets[i], book.offsets[i + 1]).filter((t) => t >= 0); + if (tokens.length === 0) { + // The cue lies wholly on unread pages: no false positive. + expect(spans[i]!.matched).toBe(false); + } else { + expect(spans[i]).toEqual({ + start: tokens[0], + end: tokens[tokens.length - 1]! + 1, + matched: true, + }); + } + }); + }), + ); + }); + + it('matches no cue when the cues share no 8-character run with the pages', () => { + // Cues from one random text, pages from another. Both use one alphabet, + // so chance runs of a few characters are common. The cue stream is + // longer than a short stream, which is compared with the pages in full. + fc.assert( + fc.property( + fc.array(word(15, 40), { minLength: 5, maxLength: 30 }), + exactBooks, + (cues, book) => { + const cueGrams = grams(cues.join('')); + fc.pre( + book.pages.every((page) => + [...grams(page.map((t) => t.text).join(''))].every((g) => !cueGrams.has(g)), + ), + ); + const spans = createAligner(cues).spans(book.pages); + expect(spans.filter((s) => s.matched)).toEqual([]); + }, + ), + ); + }); + + it('gives the same spans with the page cache as a new aligner', () => { + fc.assert( + fc.property(withReadPages(exactBooks), ([book, read]) => { + const aligner = createAligner(book.cues); + aligner.spans(readPages(book, read)); + // Read the other pages one by one, as the reader does. + const now = [...read]; + for (let p = 0; p < now.length; p++) { + if (now[p]) continue; + now[p] = true; + const pages = readPages(book, now); + expect(aligner.spans(pages)).toEqual(createAligner(book.cues).spans(pages)); + } + }), + { numRuns: 30 }, + ); + }); + + it('gives rising, in-range spans when pages repeat, overlap or come out of order', () => { + const slices = (length: number) => + fc.array( + fc.option(fc.record({ start: fc.nat({ max: length }), length: fc.nat({ max: 300 }) }), { + nil: undefined, + }), + { maxLength: 12 }, + ); + fc.assert( + fc.property( + exactBooks.chain((book) => fc.tuple(fc.constant(book), slices(book.pageOf.length))), + ([book, parts]) => { + const text = book.cues.join(''); + const pages = parts.map( + (s, p) => s && tokensOf(text.slice(s.start, s.start + s.length), p), + ); + const total = pages.reduce((n, p) => n + (p?.length ?? 0), 0); + // One character per token: the spans do not overlap. + expectRising(createAligner(book.cues).spans(pages), total, true); + }, + ), + ); + }); + + it('gives rising, in-range spans for any text', () => { fc.assert( fc.property( fc.array(fc.string({ unit: 'grapheme' }), { maxLength: 20 }), - fc.string({ unit: 'grapheme' }), - (cues, ocr) => { - const spans = alignCues(cues, ocr); + fc.array(fc.option(fc.array(fc.string({ unit: 'grapheme' })), { nil: undefined }), { + maxLength: 5, + }), + (cues, texts) => { + const pages = texts.map((page, p) => page?.map((text) => ({ text, page: p }))); + const total = pages.reduce((n, p) => n + (p?.length ?? 0), 0); + const spans = createAligner(cues).spans(pages); expect(spans).toHaveLength(cues.length); - let prevEnd = 0; - for (const s of spans) { - expect(s.start).toBeGreaterThanOrEqual(prevEnd); - expect(s.end).toBeGreaterThanOrEqual(s.start); - expect(s.end).toBeLessThanOrEqual(ocr.length); - prevEnd = s.end; - } + // A token can hold characters of two cues, so two spans can share a token. + expectRising(spans, total, false); }, ), ); }); + it('keeps 95% of the cues on read pages on their own pages under 5% noise', () => { + const noisy = books({ + cues: [100, 150], + cueLength: [10, 40], + pageLength: [100, 400], + noise: 0.05, + }); + fc.assert( + fc.property(withReadPages(noisy), ([book, read]) => { + const pages = readPages(book, read); + const tokens = pages.flatMap((p) => p ?? []); + const spans = createAligner(book.cues).spans(pages); + let onRead = 0; + let right = 0; + book.cues.forEach((_, i) => { + const own = pagesOfCue(book, i); + if (![...own].every((p) => read[p])) return; + onRead++; + const span = spans[i]!; + if (span.matched && tokens.slice(span.start, span.end).every((t) => own.has(t.page))) { + right++; + } + }); + expect(right).toBeGreaterThanOrEqual(0.95 * onRead); + }), + { numRuns: 30 }, + ); + }); +}); + +describe('alignCuesToTokens on one page', () => { + it('maps exact text to exact spans, for short cues too', () => { + fc.assert( + fc.property(fc.array(word(1, 30), { minLength: 1, maxLength: 20 }), (cues) => { + const spans = alignCuesToTokens(cues, tokensOf(cues.join(''))); + let offset = 0; + cues.forEach((cue, i) => { + expect(spans[i]).toEqual({ start: offset, end: offset + cue.length, matched: true }); + offset += cue.length; + }); + }), + ); + }); + it('keeps spans near the truth under modest random character noise', () => { // One cue with up to one edit per ten characters. const noisyCue = word(15, 40).chain((text) => @@ -73,8 +281,7 @@ describe('alignCues properties', () => { fc.property(fc.array(noisyCue, { minLength: 1, maxLength: 15 }), (items) => { const cues = items.map(([text]) => text); const noisy = items.map(([text, edits]) => applyEdits(text, edits)); - const ocr = noisy.join(''); - const spans = alignCues(cues, ocr); + const spans = alignCuesToTokens(cues, tokensOf(noisy.join(''))); let offset = 0; noisy.forEach((truth, i) => { diff --git a/test/align.test.ts b/test/align.test.ts index 0ac49c2..5e97872 100644 --- a/test/align.test.ts +++ b/test/align.test.ts @@ -1,11 +1,25 @@ +import fc from 'fast-check'; import { describe, expect, it } from 'vitest'; -import { - alignCues, - alignCuesToTokens, - buildStream, - normalizeText, - toTokenSpan, -} from '../src/align'; +import { alignCuesToTokens, buildStream, createAligner, normalizeText } from '../src/align'; +import { ENG2 } from '../fixtures/sample-text'; +import { books, char, grams, pagesOfCue, readPages, type Token } from './book'; + +/** One token for each character of `text`, all on `page`. */ +function tokensOf(text: string, page = 0): Token[] { + return [...text].map((ch) => ({ text: ch, page })); +} + +/** The opening of A Tale of Two Cities, as cues. */ +const TWO_CITIES = [ + 'It was the best of times,', + 'it was the worst of times,', + 'it was the age of wisdom,', + 'it was the age of foolishness,', + 'it was the epoch of belief,', + 'it was the epoch of incredulity,', + 'it was the season of Light,', + 'it was the season of Darkness,', +]; describe('normalizeText', () => { it('drops spaces and punctuation and folds case', () => { @@ -21,11 +35,10 @@ describe('normalizeText', () => { }); }); -describe('alignCues', () => { +describe('alignCuesToTokens on one page', () => { it('maps exact text to exact spans', () => { const cues = ['thequick', 'brownfox', 'jumps']; - const spans = alignCues(cues, cues.join('')); - expect(spans).toEqual([ + expect(alignCuesToTokens(cues, tokensOf(cues.join('')))).toEqual([ { start: 0, end: 8, matched: true }, { start: 8, end: 16, matched: true }, { start: 16, end: 21, matched: true }, @@ -35,69 +48,51 @@ describe('alignCues', () => { it('tolerates OCR substitutions and drops', () => { // "rn" read as "m", "l" read as "1", one char lost. const cues = ['themorningsun', 'wasclearandbright']; - const ocr = 'thernorningsun' + 'wasc1earandbrigt'; - const spans = alignCues(cues, ocr); + const spans = alignCuesToTokens(cues, tokensOf('thernorningsun' + 'wasc1earandbrigt')); expect(spans[0]).toEqual({ start: 0, end: 14, matched: true }); expect(spans[1]).toEqual({ start: 14, end: 30, matched: true }); }); it('skips OCR text that no cue covers, such as a page header', () => { const cues = ['firstline', 'secondline']; - const ocr = 'chapterone' + 'firstline' + 'pagenumber12' + 'secondline'; - const spans = alignCues(cues, ocr); + const spans = alignCuesToTokens( + cues, + tokensOf('chapterone' + 'firstline' + 'pagenumber12' + 'secondline'), + ); expect(spans[0]).toEqual({ start: 10, end: 19, matched: true }); expect(spans[1]).toEqual({ start: 31, end: 41, matched: true }); }); - it('gives an unmatched cue a share of the gap between its neighbors', () => { + it('leaves a cue that is not on the page unmatched, with an empty span', () => { const cues = ['abcdefgh', 'zzzzzzzz', 'ijklmnop']; - const ocr = 'abcdefgh' + 'qqqqqqqq' + 'ijklmnop'; - const spans = alignCues(cues, ocr); - expect(spans[1]).toEqual({ start: 8, end: 16, matched: false }); + const spans = alignCuesToTokens(cues, tokensOf('abcdefgh' + 'qqqqqqqq' + 'ijklmnop')); + expect(spans).toEqual([ + { start: 0, end: 8, matched: true }, + { start: 0, end: 0, matched: false }, + { start: 16, end: 24, matched: true }, + ]); }); it('aligns Japanese text at character level with OCR noise', () => { const cues = ['吾輩は猫である', '名前はまだ無い', 'どこで生れたか']; // 猫 read as 描, one char doubled, one char lost. - const ocr = '吾輩は描である' + '名前はまだだ無い' + 'どこで生れか'; - const spans = alignCues(cues, ocr); + const spans = alignCuesToTokens( + cues, + tokensOf('吾輩は描である' + '名前はまだだ無い' + 'どこで生れか'), + ); expect(spans[0]).toEqual({ start: 0, end: 7, matched: true }); expect(spans[1]).toEqual({ start: 7, end: 15, matched: true }); expect(spans[2]).toEqual({ start: 15, end: 21, matched: true }); }); it('returns empty spans for empty input', () => { - expect(alignCues([], 'abc')).toEqual([]); - expect(alignCues(['abc'], '')).toEqual([{ start: 0, end: 0, matched: false }]); + expect(alignCuesToTokens([], tokensOf('abc'))).toEqual([]); + expect(alignCuesToTokens(['abc'], [])).toEqual([{ start: 0, end: 0, matched: false }]); }); -}); -describe('buildStream and toTokenSpan', () => { - it('maps chars back to the tokens that hold them', () => { - const stream = buildStream(['The', 'quick,', '—', 'fox']); - expect(stream.chars).toBe('thequickfox'); - expect(stream.tokenOf).toEqual([0, 0, 0, 1, 1, 1, 1, 1, 3, 3, 3]); - expect(toTokenSpan({ start: 3, end: 11, matched: true }, stream)).toEqual({ - start: 1, - end: 4, - matched: true, - }); - }); - - it('points an empty span at the token at that position', () => { - const stream = buildStream(['ab', 'cd']); - expect(toTokenSpan({ start: 2, end: 2, matched: false }, stream)).toEqual({ - start: 1, - end: 1, - matched: false, - }); - }); -}); - -describe('alignCuesToTokens', () => { it('runs the full pipeline on word tokens with punctuation tokens', () => { const tokens = ['It', 'was', 'a', 'dark', 'and', 'stormy', 'night', '.', 'The', 'rain'].map( - (text) => ({ text }), + (text) => ({ text, page: 0 }), ); const spans = alignCuesToTokens(['It was a dark', 'and stormy night.', 'The rain'], tokens); expect(spans).toEqual([ @@ -108,10 +103,128 @@ describe('alignCuesToTokens', () => { }); it('works with one token per character, as for vertical Japanese', () => { - const text = '春はあけぼの。やうやう白くなりゆく山ぎは'; - const tokens = [...text].map((ch) => ({ text: ch })); - const spans = alignCuesToTokens(['春はあけぼの', 'やうやう白くなりゆく山ぎは'], tokens); + const spans = alignCuesToTokens( + ['春はあけぼの', 'やうやう白くなりゆく山ぎは'], + tokensOf('春はあけぼの。やうやう白くなりゆく山ぎは'), + ); expect(spans[0]).toEqual({ start: 0, end: 6, matched: true }); expect(spans[1]).toEqual({ start: 7, end: 20, matched: true }); }); + + it('refuses tokens that are not in page order', () => { + const tokens = [ + { text: 'a', page: 1 }, + { text: 'b', page: 0 }, + ]; + expect(() => alignCuesToTokens(['ab'], tokens)).toThrow(/page order/); + }); +}); + +describe('buildStream', () => { + it('maps chars back to the tokens that hold them', () => { + const stream = buildStream(['The', 'quick,', '—', 'fox']); + expect(stream.chars).toBe('thequickfox'); + expect(stream.tokenOf).toEqual([0, 0, 0, 1, 1, 1, 1, 1, 3, 3, 3]); + }); +}); + +describe('createAligner', () => { + const text = normalizeText(TWO_CITIES.join('')); + const cues = text.match(/.{1,20}/g)!; + + it('indexes the tokens of the read pages, joined in page order', () => { + const pages = [tokensOf(text.slice(0, 60), 0), undefined, tokensOf(text.slice(60), 2)]; + const spans = createAligner(cues).spans(pages); + expect(spans.map((s) => s.matched)).toEqual(cues.map(() => true)); + expect(spans[0]).toEqual({ start: 0, end: 20, matched: true }); + expect(spans[3]).toEqual({ start: 60, end: 80, matched: true }); + }); + + it('leaves the cues of an unread page unmatched', () => { + const pages = [undefined, undefined, tokensOf(text.slice(60), 2)]; + const spans = createAligner(cues).spans(pages); + expect(spans.slice(0, 3).map((s) => s.matched)).toEqual([false, false, false]); + expect(spans[3]).toEqual({ start: 0, end: 20, matched: true }); + }); + + it('matches no cue of subtitles for another text (regression)', () => { + // The opening of A Tale of Two Cities over the two pages of the + // sample-eng-2p fixture. The texts share no 8-character run, so no + // page has an anchor. One diff of all cues matched every cue on chance + // runs of two characters. + const pages = ENG2.pages.map((lines, p) => tokensOf(lines.join(' '), p)); + const pageGrams = grams(ENG2.pages.flat().join('')); + expect([...grams(TWO_CITIES.join(''))].filter((g) => pageGrams.has(g))).toEqual([]); + + const spans = alignCuesToTokens(TWO_CITIES, pages.flat()); + expect(spans.filter((s) => s.matched)).toEqual([]); + }); + + it('matches no cue to a wrong page when only 3 pages in the middle are read (regression)', () => { + // The measured defect in small form: a book of 20 pages with 5% OCR + // noise, and only pages 9 to 11 read. One diff of all cues against + // those pages matched cues to wrong places. + const [book] = fc.sample( + books({ cues: [220, 220], cueLength: [15, 40], pageLength: [300, 300], noise: 0.05 }), + { seed: 3, numRuns: 1 }, + ); + const read = book!.pages.map((_, p) => p >= 9 && p <= 11); + const pages = readPages(book!, read); + const tokens = pages.flatMap((p) => p ?? []); + const spans = alignCuesToTokens(book!.cues, tokens); + + const wrong: number[] = []; + const missed: number[] = []; + book!.cues.forEach((_, i) => { + const own = pagesOfCue(book!, i); + const span = spans[i]!; + const onOwnPages = tokens.slice(span.start, span.end).every((t) => own.has(t.page)); + if (span.matched && !onOwnPages) wrong.push(i); + if ([...own].every((p) => read[p]) && !(span.matched && onOwnPages)) missed.push(i); + }); + expect(wrong).toEqual([]); + expect(missed).toEqual([]); + }); +}); + +describe('createAligner speed', () => { + it('aligns a book of 300 000 characters in under 3 s, and one more page in under 100 ms', () => { + const LENGTH = 300_000; + const PAGE = 600; + // 5% noise: 3 codes in 60 are a drop, a substitution and an insertion. + const codes = [...'k'.repeat(57), 'd', 's', 'i']; + const [sample] = fc.sample( + fc.tuple( + fc.string({ unit: char, minLength: LENGTH, maxLength: LENGTH }), + fc.string({ unit: fc.constantFrom(...codes), minLength: LENGTH, maxLength: LENGTH }), + ), + { seed: 1, numRuns: 1 }, + ); + const [text, noise] = sample!; + const cues = text.match(/.{1,40}/g)!; + const pages: Token[][] = []; + for (let i = 0; i < LENGTH; i++) { + const page = Math.floor(i / PAGE); + const tokens = (pages[page] ??= []); + const ch = text[i]!; + const other = text[(i + 1) % LENGTH]!; + if (noise[i] === 's') tokens.push({ text: other, page }); + else if (noise[i] === 'i') tokens.push({ text: other, page }, { text: ch, page }); + else if (noise[i] !== 'd') tokens.push({ text: ch, page }); + } + + let t = performance.now(); + const aligner = createAligner(cues); + const spans = aligner.spans(pages); + const all = performance.now() - t; + // Read one page again: a new token array, as after a change of the OCR language. + pages[250] = [...pages[250]!]; + t = performance.now(); + aligner.spans(pages); + const one = performance.now() - t; + + expect(spans.filter((s) => s.matched).length).toBeGreaterThanOrEqual(0.99 * cues.length); + expect(all).toBeLessThan(3000); + expect(one).toBeLessThan(100); + }, 30_000); }); diff --git a/test/book.ts b/test/book.ts new file mode 100644 index 0000000..221f6ec --- /dev/null +++ b/test/book.ts @@ -0,0 +1,114 @@ +// Synthetic books for the alignment tests: cue texts, pages of one token +// per character, and OCR noise. All randomness comes from fast-check. + +import fc from 'fast-check'; +import { normalizeText, type OcrToken } from '../src/align'; + +/** A small mixed alphabet: Latin, kana and kanji, so CJK is always covered. */ +export const ALPHABET = [...'abcdefghijklmnop', ...'あいうえおかきくけこ', ...'吾輩猫名前春夏秋冬']; + +export const char = fc.constantFrom(...ALPHABET); + +export const word = (min: number, max: number): fc.Arbitrary<string> => + fc.array(char, { minLength: min, maxLength: max }).map((cs) => cs.join('')); + +export type Token = Pick<OcrToken, 'text' | 'page'>; + +export interface Book { + cues: string[]; + /** The start of each cue in the book text, and the end of the text last. */ + offsets: number[]; + /** The page of each character of the book text. */ + pageOf: number[]; + /** The tokens of each page, one for each character, after the noise. */ + pages: Token[][]; +} + +/** + * One noise draw for each character: a number in [0, 1) and a character + * for a substitution or an insertion. No bias, so the rate is exact. + */ +const draw = fc.tuple( + fc.noBias(fc.integer({ min: 0, max: 9999 })).map((n) => n / 10000), + char, +); + +export interface BookShape { + /** The count of cues. */ + cues: [number, number]; + /** The length of a cue in characters. */ + cueLength: [number, number]; + /** The length of a page in characters. The last page takes the rest. */ + pageLength: [number, number]; + /** The share of characters with an edit: a drop, a substitution or an insertion. */ + noise: number; +} + +/** Books of the given shape. */ +export function books(shape: BookShape): fc.Arbitrary<Book> { + return fc + .record({ + cues: fc.array(word(shape.cueLength[0], shape.cueLength[1]), { + minLength: shape.cues[0], + maxLength: shape.cues[1], + }), + pageLength: fc.integer({ min: shape.pageLength[0], max: shape.pageLength[1] }), + }) + .chain(({ cues, pageLength }) => { + const length = shape.noise > 0 ? cues.join('').length : 0; + return fc + .array(draw, { minLength: length, maxLength: length }) + .map((draws) => makeBook(cues, pageLength, draws, shape.noise)); + }); +} + +function makeBook( + cues: string[], + pageLength: number, + draws: readonly [number, string][], + noise: number, +): Book { + const text = cues.join(''); + const offsets: number[] = []; + let at = 0; + for (const cue of cues) { + offsets.push(at); + at += cue.length; + } + offsets.push(at); + const pageCount = Math.max(1, Math.floor(text.length / pageLength)); + const pageOf = [...text].map((_, i) => Math.min(pageCount - 1, Math.floor(i / pageLength))); + const pages: Token[][] = Array.from({ length: pageCount }, () => []); + [...text].forEach((ch, i) => { + const page = pageOf[i]!; + const [r, other] = draws[i] ?? [1, ch]; + // A third of the edits drop the character, a third put another in its + // place, and a third put another before it. + let read = [ch]; + if (r < noise / 3) read = []; + else if (r < (2 * noise) / 3) read = [other]; + else if (r < noise) read = [other, ch]; + for (const text of read) pages[page]!.push({ text, page }); + }); + return { cues, offsets, pageOf, pages }; +} + +/** The pages that hold the characters of cue `i`. */ +export function pagesOfCue(book: Book, i: number): Set<number> { + const pages = new Set<number>(); + for (let k = book.offsets[i]!; k < book.offsets[i + 1]!; k++) pages.add(book.pageOf[k]!); + return pages; +} + +/** The runs of `k` characters in the normalized form of `text`. */ +export function grams(text: string, k = 8): Set<string> { + const norm = normalizeText(text); + const out = new Set<string>(); + for (let i = 0; i + k <= norm.length; i++) out.add(norm.slice(i, i + k)); + return out; +} + +/** The read pages of a book: `read[p]` true keeps page p, else it is undefined. */ +export function readPages(book: Book, read: readonly boolean[]): (Token[] | undefined)[] { + return book.pages.map((tokens, p) => (read[p] ? tokens : undefined)); +}