From 81bd355bd1d225f485e5a84f569a996db0f2d9ec Mon Sep 17 00:00:00 2001 From: Richard R Date: Sat, 6 Jun 2026 10:45:19 -0600 Subject: [PATCH] refactor(highlight): share multilingual token alignment --- src/hooks/epub/useEPUBHighlighting.ts | 30 ++- src/lib/client/epub/epub-word-highlight.ts | 69 +----- src/lib/client/highlight-token-alignment.ts | 232 ++++++++++++++++++ src/lib/client/html/highlight.ts | 187 ++------------ src/lib/client/pdf-highlight-match.ts | 101 -------- src/lib/client/pdf-highlight-worker.ts | 7 +- src/lib/client/pdf.ts | 203 ++++----------- tests/unit/epub-word-highlight.vitest.spec.ts | 29 --- .../highlight-token-alignment.vitest.spec.ts | 92 +++++++ tests/unit/pdf-highlight-match.vitest.spec.ts | 29 --- 10 files changed, 432 insertions(+), 547 deletions(-) create mode 100644 src/lib/client/highlight-token-alignment.ts delete mode 100644 src/lib/client/pdf-highlight-match.ts create mode 100644 tests/unit/highlight-token-alignment.vitest.spec.ts delete mode 100644 tests/unit/pdf-highlight-match.vitest.spec.ts diff --git a/src/hooks/epub/useEPUBHighlighting.ts b/src/hooks/epub/useEPUBHighlighting.ts index 83beafa..7362d3e 100644 --- a/src/hooks/epub/useEPUBHighlighting.ts +++ b/src/hooks/epub/useEPUBHighlighting.ts @@ -4,12 +4,15 @@ import { useCallback, useEffect, type MutableRefObject, type RefObject } from 'r import type { Rendition } from 'epubjs'; import { - buildMonotonicWordToTokenMap, buildWordHighlightCacheKey, resolveAlignmentWordSourceRange, tokenizeCanonicalSegment, type EpubCanonicalWordToken, } from '@/lib/client/epub/epub-word-highlight'; +import { + buildAlignmentTokenRanges, + type HighlightTokenRange, +} from '@/lib/client/highlight-token-alignment'; import { createRangeFromMappedOffsets, resolveVisibleSegmentRange, @@ -20,7 +23,7 @@ import type { TTSSentenceAlignment } from '@/types/tts'; export type EpubWordHighlightMapCache = { key: string; - wordToToken: number[]; + wordToTokenRange: Array; tokens: EpubCanonicalWordToken[]; }; @@ -159,19 +162,28 @@ export function useEPUBHighlighting({ wordHighlightMapCacheRef.current = { key: cacheKey, tokens, - wordToToken: buildMonotonicWordToTokenMap(words, tokens), + wordToTokenRange: buildAlignmentTokenRanges( + words, + tokens.map((token) => token.norm), + { minimumSimilarity: 0.8 }, + ), }; } const cached = wordHighlightMapCacheRef.current; - const tokenIndex = cached.wordToToken[wordIndex] ?? -1; - if (tokenIndex < 0) return; + const tokenRange = cached.wordToTokenRange[wordIndex]; + if (!tokenRange) return; - const token = cached.tokens[tokenIndex]; - if (!token) return; - if (token.sourceStart < resolved.startOffset || token.sourceEnd > resolved.endOffset) return; + const firstToken = cached.tokens[tokenRange.start]; + const lastToken = cached.tokens[tokenRange.end]; + if (!firstToken || !lastToken) return; + if (firstToken.sourceStart < resolved.startOffset || lastToken.sourceEnd > resolved.endOffset) return; - const wordRange = createRangeFromMappedOffsets(resolved.map, token.sourceStart, token.sourceEnd); + const wordRange = createRangeFromMappedOffsets( + resolved.map, + firstToken.sourceStart, + lastToken.sourceEnd, + ); if (!wordRange) return; try { diff --git a/src/lib/client/epub/epub-word-highlight.ts b/src/lib/client/epub/epub-word-highlight.ts index 7007ce5..de44021 100644 --- a/src/lib/client/epub/epub-word-highlight.ts +++ b/src/lib/client/epub/epub-word-highlight.ts @@ -1,6 +1,7 @@ import type { CanonicalTtsSegment } from '@/lib/shared/tts-segment-plan'; import type { TTSSentenceAlignment, TTSSentenceWord } from '@/types/tts'; -import { normalizeUnicodeToken, segmentWords } from '@/lib/shared/language'; +import { segmentWords } from '@/lib/shared/language'; +import { normalizeHighlightToken } from '@/lib/client/highlight-token-alignment'; export type EpubCanonicalWordToken = { norm: string; @@ -9,7 +10,7 @@ export type EpubCanonicalWordToken = { }; export const normalizeWordForHighlight = (text: string): string => - normalizeUnicodeToken(text); + normalizeHighlightToken(text); export const resolveAlignmentWordSourceRange = ( segment: CanonicalTtsSegment, @@ -37,70 +38,6 @@ export const tokenizeCanonicalSegment = ( })) .filter((token) => Boolean(token.norm)); -export const buildMonotonicWordToTokenMap = ( - alignmentWords: TTSSentenceAlignment['words'], - segmentTokens: EpubCanonicalWordToken[], -): number[] => { - const alignmentTokens = alignmentWords.map((word) => normalizeWordForHighlight(word.text)); - const wordToToken = new Array(alignmentWords.length).fill(-1); - const m = alignmentTokens.length; - const n = segmentTokens.length; - if (!m || !n) return wordToToken; - - const dp: number[][] = Array.from({ length: m + 1 }, () => new Array(n + 1).fill(0)); - const bt: number[][] = Array.from({ length: m + 1 }, () => new Array(n + 1).fill(0)); - - for (let i = 1; i <= m; i += 1) { - for (let j = 1; j <= n; j += 1) { - let best = dp[i - 1][j - 1]; - let move = 0; - - const alignmentNorm = alignmentTokens[i - 1]; - const segmentNorm = segmentTokens[j - 1].norm; - if (alignmentNorm && alignmentNorm === segmentNorm) { - const positionPenalty = - m <= 1 || n <= 1 - ? 0 - : Math.abs((i - 1) / (m - 1) - (j - 1) / (n - 1)); - best = dp[i - 1][j - 1] + 10 - positionPenalty; - move = 1; - } - - if (dp[i - 1][j] > best) { - best = dp[i - 1][j]; - move = 2; - } - if (dp[i][j - 1] > best) { - best = dp[i][j - 1]; - move = 3; - } - - dp[i][j] = best; - bt[i][j] = move; - } - } - - let i = m; - let j = n; - while (i > 0 && j > 0) { - const move = bt[i][j]; - if (move === 1) { - wordToToken[i - 1] = j - 1; - i -= 1; - j -= 1; - } else if (move === 2) { - i -= 1; - } else if (move === 3) { - j -= 1; - } else { - i -= 1; - j -= 1; - } - } - - return wordToToken; -}; - export const buildWordHighlightCacheKey = ( segment: CanonicalTtsSegment, alignment: TTSSentenceAlignment, diff --git a/src/lib/client/highlight-token-alignment.ts b/src/lib/client/highlight-token-alignment.ts new file mode 100644 index 0000000..26210ad --- /dev/null +++ b/src/lib/client/highlight-token-alignment.ts @@ -0,0 +1,232 @@ +import { CmpStr } from 'cmpstr'; + +import { normalizeUnicodeToken } from '@/lib/shared/language'; +import type { TTSSentenceAlignment } from '@/types/tts'; + +const cmp = CmpStr.create().setMetric('dice').setFlags('itw'); + +export interface HighlightTokenRange { + start: number; + end: number; +} + +export interface HighlightTokenMatchResult extends HighlightTokenRange { + rating: number; + lengthDiff: number; +} + +export function normalizeHighlightToken(text: string): string { + return normalizeUnicodeToken(text); +} + +export function findBestHighlightTokenMatch( + patternTokens: string[], + tokenTexts: string[], +): HighlightTokenMatchResult { + const normalizedPattern = patternTokens.map(normalizeHighlightToken).filter(Boolean); + const normalizedTargets = tokenTexts.map(normalizeHighlightToken); + const cleanPattern = normalizedPattern.join(''); + const patternLen = cleanPattern.length; + const responseBase: HighlightTokenMatchResult = { + start: -1, + end: -1, + rating: 0, + lengthDiff: Number.POSITIVE_INFINITY, + }; + + if (!patternLen || !normalizedTargets.length) return responseBase; + + const patternTokenCount = normalizedPattern.length; + const minWindowTokens = Math.max(1, Math.floor(patternTokenCount * 0.6)); + const maxWindowTokens = Math.max( + minWindowTokens, + Math.ceil(patternTokenCount * 1.4), + ); + + let bestStart = -1; + let bestEnd = -1; + let bestRating = 0; + let bestLengthDiff = Number.POSITIVE_INFINITY; + + for (let start = 0; start < normalizedTargets.length; start += 1) { + let combined = ''; + + for ( + let offset = 0; + offset < maxWindowTokens && start + offset < normalizedTargets.length; + offset += 1 + ) { + combined += normalizedTargets[start + offset]; + + const windowSize = offset + 1; + if (windowSize < minWindowTokens) continue; + if (combined.length > patternLen * 2) break; + + const similarity = cmp.compare(combined, cleanPattern); + const lengthDiff = Math.abs(combined.length - patternLen); + const lengthPenalty = lengthDiff / patternLen; + const adjustedRating = similarity * (1 - lengthPenalty * 0.3); + + let boostedRating = adjustedRating; + const maxPrefixCheck = Math.min(windowSize, normalizedPattern.length, 5); + let prefixMatches = 0; + for (let i = 0; i < maxPrefixCheck; i += 1) { + const tokenSim = cmp.compare(normalizedTargets[start + i], normalizedPattern[i]); + if (tokenSim < 0.8) break; + prefixMatches += 1; + } + + if (prefixMatches > 0) { + boostedRating = adjustedRating * (1 + (prefixMatches / maxPrefixCheck) * 0.25); + } + + if ( + boostedRating > bestRating + || ( + Math.abs(boostedRating - bestRating) < 1e-3 + && lengthDiff < bestLengthDiff + ) + ) { + bestRating = boostedRating; + bestLengthDiff = lengthDiff; + bestStart = start; + bestEnd = start + offset; + } + } + } + + return { + start: bestStart, + end: bestEnd, + rating: bestRating, + lengthDiff: bestLengthDiff, + }; +} + +function buildExactConcatenatedRanges( + alignmentNorms: string[], + targetNorms: string[], +): Array | null { + if (alignmentNorms.join('') !== targetNorms.join('')) return null; + + const targetOffsets: HighlightTokenRange[] = []; + let targetCursor = 0; + for (const norm of targetNorms) { + targetOffsets.push({ start: targetCursor, end: targetCursor + norm.length }); + targetCursor += norm.length; + } + + const ranges: Array = []; + let alignmentCursor = 0; + for (const norm of alignmentNorms) { + const alignmentStart = alignmentCursor; + const alignmentEnd = alignmentStart + norm.length; + alignmentCursor = alignmentEnd; + + let first = -1; + let last = -1; + for (let i = 0; i < targetOffsets.length; i += 1) { + const target = targetOffsets[i]; + if (target.end <= alignmentStart) continue; + if (target.start >= alignmentEnd) break; + if (first === -1) first = i; + last = i; + } + ranges.push(first === -1 ? null : { start: first, end: last }); + } + + return ranges; +} + +export function buildAlignmentTokenRanges( + alignmentWords: TTSSentenceAlignment['words'], + targetTexts: string[], + options: { fillGaps?: boolean; minimumSimilarity?: number } = {}, +): Array { + const alignmentNorms = alignmentWords.map((word) => normalizeHighlightToken(word.text)); + const targetNorms = targetTexts.map(normalizeHighlightToken); + const ranges = new Array(alignmentWords.length).fill(null); + + const exactRanges = buildExactConcatenatedRanges(alignmentNorms, targetNorms); + if (exactRanges) return exactRanges; + + const targets = targetNorms + .map((norm, tokenIndex) => ({ norm, tokenIndex })) + .filter((token) => Boolean(token.norm)); + const alignments = alignmentNorms + .map((norm, wordIndex) => ({ norm, wordIndex })) + .filter((word) => Boolean(word.norm)); + const m = targets.length; + const n = alignments.length; + if (!m || !n) return ranges; + + const dp: number[][] = Array.from({ length: m + 1 }, () => + new Array(n + 1).fill(Number.POSITIVE_INFINITY), + ); + const bt: number[][] = Array.from({ length: m + 1 }, () => + new Array(n + 1).fill(0), + ); + dp[0][0] = 0; + const gapCost = 0.7; + + for (let i = 0; i <= m; i += 1) { + for (let j = 0; j <= n; j += 1) { + if (i > 0 && j > 0) { + const a = targets[i - 1].norm; + const b = alignments[j - 1].norm; + const similarity = a === b ? 1 : cmp.compare(a, b); + const candidate = dp[i - 1][j - 1] + (1 - similarity); + if (candidate < dp[i][j]) { + dp[i][j] = candidate; + bt[i][j] = 0; + } + } + if (i > 0 && dp[i - 1][j] + gapCost < dp[i][j]) { + dp[i][j] = dp[i - 1][j] + gapCost; + bt[i][j] = 1; + } + if (j > 0 && dp[i][j - 1] + gapCost < dp[i][j]) { + dp[i][j] = dp[i][j - 1] + gapCost; + bt[i][j] = 2; + } + } + } + + let i = m; + let j = n; + while (i > 0 || j > 0) { + const move = bt[i][j]; + if (i > 0 && j > 0 && move === 0) { + const tokenIndex = targets[i - 1].tokenIndex; + const a = targets[i - 1].norm; + const b = alignments[j - 1].norm; + const similarity = a === b ? 1 : cmp.compare(a, b); + if (similarity >= (options.minimumSimilarity ?? 0)) { + ranges[alignments[j - 1].wordIndex] = { start: tokenIndex, end: tokenIndex }; + } + i -= 1; + j -= 1; + } else if (i > 0 && (move === 1 || j === 0)) { + i -= 1; + } else if (j > 0 && (move === 2 || i === 0)) { + j -= 1; + } else { + break; + } + } + + if (!options.fillGaps) return ranges; + + let lastSeen: HighlightTokenRange | null = null; + for (let k = 0; k < ranges.length; k += 1) { + if (ranges[k]) lastSeen = ranges[k]; + else if (lastSeen) ranges[k] = lastSeen; + } + let nextSeen: HighlightTokenRange | null = null; + for (let k = ranges.length - 1; k >= 0; k -= 1) { + if (ranges[k]) nextSeen = ranges[k]; + else if (nextSeen) ranges[k] = nextSeen; + } + + return ranges; +} diff --git a/src/lib/client/html/highlight.ts b/src/lib/client/html/highlight.ts index db26cc2..a78a70a 100644 --- a/src/lib/client/html/highlight.ts +++ b/src/lib/client/html/highlight.ts @@ -11,14 +11,18 @@ * sentence * - `WORD` — saturated background on the currently-spoken word * - * Word-to-DOM alignment is done via Needleman-Wunsch (same approach the PDF - * reader uses) so DOM token counts that diverge from whisper's word count - * still produce a smooth, monotonic word highlight rather than a proportional - * approximation that snaps around when the counts disagree. + * Word-to-DOM alignment uses the shared viewer token-range mapper so DOM token + * counts that diverge from the timed alignment still produce a smooth, + * monotonic highlight across languages. */ -import { CmpStr } from 'cmpstr'; import type { TTSSentenceAlignment } from '@/types/tts'; -import { normalizeUnicodeToken, segmentWords } from '@/lib/shared/language'; +import { segmentWords } from '@/lib/shared/language'; +import { + buildAlignmentTokenRanges, + findBestHighlightTokenMatch, + normalizeHighlightToken, + type HighlightTokenRange, +} from '@/lib/client/highlight-token-alignment'; export const HTML_SENTENCE_CLASS = 'openreader-html-highlight-sentence'; export const HTML_WORD_CLASS = 'openreader-html-highlight-word'; @@ -30,8 +34,6 @@ interface DomToken { norm: string; } -const cmp = CmpStr.create().setMetric('dice').setFlags('itw'); - let sentenceWraps: HTMLSpanElement[] = []; let wordWraps: HTMLSpanElement[] = []; @@ -46,15 +48,15 @@ interface SentenceState { // is in place. Stable across word wrap/unwrap cycles because clear() calls // `parent.normalize()` which restores the original text-node structure. wordTokens: DomToken[]; - // For an alignment we've already seen, the cached wordIndex → tokenIndex map. + // For an alignment we've already seen, the cached wordIndex → token range map. alignment: TTSSentenceAlignment | null; - wordToToken: number[] | null; + wordToTokenRange: Array | null; } let sentenceState: SentenceState | null = null; function normalizeWord(word: string): string { - return normalizeUnicodeToken(word); + return normalizeHighlightToken(word); } function tokenizePattern(pattern: string, language?: string): string[] { @@ -171,43 +173,9 @@ function collectTokensInsideWraps(wraps: HTMLSpanElement[], language?: string): } function findBestWindow(tokens: DomToken[], patternTokens: string[]): { start: number; end: number } | null { - if (!tokens.length || !patternTokens.length) return null; - const pLen = patternTokens.length; - - let bestStart = -1; - let bestEnd = -1; - let bestScore = 0; - - for (let i = 0; i + Math.max(1, Math.ceil(pLen * 0.5)) - 1 < tokens.length; i += 1) { - if (tokens[i].norm !== patternTokens[0]) continue; - let matches = 1; - let domCursor = i + 1; - for (let p = 1; p < pLen && domCursor < tokens.length; p += 1) { - let stepped = false; - for (let k = 0; k < 3 && domCursor + k < tokens.length; k += 1) { - if (tokens[domCursor + k].norm === patternTokens[p]) { - matches += 1; - domCursor += k + 1; - stepped = true; - break; - } - } - if (!stepped) domCursor += 1; - } - const end = Math.min(tokens.length - 1, domCursor - 1); - const score = matches / pLen; - if (score > bestScore) { - bestScore = score; - bestStart = i; - bestEnd = end; - if (score >= 0.95) break; - } - } - - if (bestScore >= 0.5 && bestStart !== -1) { - return { start: bestStart, end: bestEnd }; - } - return null; + const match = findBestHighlightTokenMatch(patternTokens, tokens.map((token) => token.norm)); + if (match.start === -1 || match.rating < 0.5) return null; + return { start: match.start, end: match.end }; } function wrapTokenRange(tokens: DomToken[], start: number, end: number, className: string): HTMLSpanElement[] { @@ -274,116 +242,11 @@ export function highlightHtmlSentence( sentence, wordTokens: collectTokensInsideWraps(sentenceWraps, language), alignment: null, - wordToToken: null, + wordToTokenRange: null, }; return true; } -/** - * Build a wordIndex → tokenIndex map via Needleman-Wunsch alignment between - * whisper's word list and the DOM tokens inside the sentence. Mirrors the - * approach in `src/lib/client/pdf.ts#highlightWordIndex` so PDF and HTML - * highlights behave the same way under count mismatches (contractions, - * stripped punctuation, missing whitespace, etc.). - */ -function buildAlignmentMap( - alignment: TTSSentenceAlignment, - domTokens: DomToken[], -): number[] { - const words = alignment.words || []; - const wordToToken = new Array(words.length).fill(-1); - - const domFiltered: { tokenIndex: number; norm: string }[] = []; - for (let i = 0; i < domTokens.length; i += 1) { - const norm = domTokens[i].norm; - if (norm) domFiltered.push({ tokenIndex: i, norm }); - } - - const ttsFiltered: { wordIndex: number; norm: string }[] = []; - for (let i = 0; i < words.length; i += 1) { - const norm = normalizeWord(words[i].text); - if (norm) ttsFiltered.push({ wordIndex: i, norm }); - } - - const m = domFiltered.length; - const n = ttsFiltered.length; - if (!m || !n) return wordToToken; - - const dp: number[][] = Array.from({ length: m + 1 }, () => - new Array(n + 1).fill(Number.POSITIVE_INFINITY), - ); - const bt: number[][] = Array.from({ length: m + 1 }, () => - new Array(n + 1).fill(0), - ); // 0=diag (substitute), 1=up (skip dom), 2=left (skip tts) - - dp[0][0] = 0; - const GAP_COST = 0.7; - - for (let i = 0; i <= m; i += 1) { - for (let j = 0; j <= n; j += 1) { - if (i > 0 && j > 0) { - const a = domFiltered[i - 1].norm; - const b = ttsFiltered[j - 1].norm; - const sim = a === b ? 1 : cmp.compare(a, b); - const cand = dp[i - 1][j - 1] + (1 - sim); - if (cand < dp[i][j]) { - dp[i][j] = cand; - bt[i][j] = 0; - } - } - if (i > 0) { - const cand = dp[i - 1][j] + GAP_COST; - if (cand < dp[i][j]) { - dp[i][j] = cand; - bt[i][j] = 1; - } - } - if (j > 0) { - const cand = dp[i][j - 1] + GAP_COST; - if (cand < dp[i][j]) { - dp[i][j] = cand; - bt[i][j] = 2; - } - } - } - } - - let i = m; - let j = n; - while (i > 0 || j > 0) { - const move = bt[i][j]; - if (i > 0 && j > 0 && move === 0) { - const domIdx = domFiltered[i - 1].tokenIndex; - const ttsIdx = ttsFiltered[j - 1].wordIndex; - if (wordToToken[ttsIdx] === -1) wordToToken[ttsIdx] = domIdx; - i -= 1; - j -= 1; - } else if (i > 0 && (move === 1 || j === 0)) { - i -= 1; - } else if (j > 0 && (move === 2 || i === 0)) { - j -= 1; - } else { - break; - } - } - - // Forward-fill, then backward-fill, so every wordIndex has a nearest known - // DOM token. This keeps the word highlight stable when whisper emits a - // word that didn't survive normalization (e.g. an apostrophe-only token). - let lastSeen = -1; - for (let k = 0; k < wordToToken.length; k += 1) { - if (wordToToken[k] !== -1) lastSeen = wordToToken[k]; - else if (lastSeen !== -1) wordToToken[k] = lastSeen; - } - let nextSeen = -1; - for (let k = wordToToken.length - 1; k >= 0; k -= 1) { - if (wordToToken[k] !== -1) nextSeen = wordToToken[k]; - else if (nextSeen !== -1) wordToToken[k] = nextSeen; - } - - return wordToToken; -} - export function highlightHtmlWord( container: HTMLElement | null | undefined, alignment: TTSSentenceAlignment | undefined, @@ -403,16 +266,20 @@ export function highlightHtmlWord( if (!words.length || wordIndex >= words.length) return false; // (Re)build the alignment map when this is a new alignment object. - if (sentenceState.alignment !== alignment || !sentenceState.wordToToken) { + if (sentenceState.alignment !== alignment || !sentenceState.wordToTokenRange) { sentenceState.alignment = alignment; - sentenceState.wordToToken = buildAlignmentMap(alignment, sentenceState.wordTokens); + sentenceState.wordToTokenRange = buildAlignmentTokenRanges( + alignment.words, + sentenceState.wordTokens.map((token) => token.norm), + { fillGaps: true }, + ); } - const tokenIndex = sentenceState.wordToToken[wordIndex]; - if (tokenIndex === undefined || tokenIndex < 0) return false; - if (tokenIndex >= sentenceState.wordTokens.length) return false; + const tokenRange = sentenceState.wordToTokenRange[wordIndex]; + if (!tokenRange) return false; + if (tokenRange.start < 0 || tokenRange.end >= sentenceState.wordTokens.length) return false; - wordWraps = wrapTokenRange(sentenceState.wordTokens, tokenIndex, tokenIndex, HTML_WORD_CLASS); + wordWraps = wrapTokenRange(sentenceState.wordTokens, tokenRange.start, tokenRange.end, HTML_WORD_CLASS); return wordWraps.length > 0; } diff --git a/src/lib/client/pdf-highlight-match.ts b/src/lib/client/pdf-highlight-match.ts deleted file mode 100644 index b16ef5d..0000000 --- a/src/lib/client/pdf-highlight-match.ts +++ /dev/null @@ -1,101 +0,0 @@ -import { CmpStr } from 'cmpstr'; - -const cmp = CmpStr.create().setMetric('dice').setFlags('itw'); - -export interface HighlightTokenMatchResult { - bestStart: number; - bestEnd: number; - rating: number; - lengthDiff: number; -} - -export function findBestHighlightTokenMatch( - patternTokens: string[], - tokenTexts: string[], -): HighlightTokenMatchResult { - const cleanPatternTokens = patternTokens.map((token) => token.trim()).filter(Boolean); - const cleanPattern = cleanPatternTokens.join(''); - const patternLen = cleanPattern.length; - const responseBase: HighlightTokenMatchResult = { - bestStart: -1, - bestEnd: -1, - rating: 0, - lengthDiff: Number.POSITIVE_INFINITY, - }; - - if (!patternLen || !tokenTexts.length) return responseBase; - - const patternTokenCount = cleanPatternTokens.length; - const minWindowTokens = Math.max(1, Math.floor(patternTokenCount * 0.6)); - const maxWindowTokens = Math.max( - minWindowTokens, - Math.ceil(patternTokenCount * 1.4), - ); - - let bestStart = -1; - let bestEnd = -1; - let bestRating = 0; - let bestLengthDiff = Number.POSITIVE_INFINITY; - - for (let start = 0; start < tokenTexts.length; start += 1) { - let combined = ''; - - for ( - let offset = 0; - offset < maxWindowTokens && start + offset < tokenTexts.length; - offset += 1 - ) { - const token = tokenTexts[start + offset]; - combined += token; - - const windowSize = offset + 1; - if (windowSize < minWindowTokens) continue; - if (combined.length > patternLen * 2) break; - - const similarity = cmp.compare(combined, cleanPattern); - const lengthDiff = Math.abs(combined.length - patternLen); - const lengthPenalty = lengthDiff / patternLen; - const adjustedRating = similarity * (1 - lengthPenalty * 0.3); - - let boostedRating = adjustedRating; - const windowTokens = tokenTexts.slice(start, start + windowSize); - const maxPrefixCheck = Math.min( - windowTokens.length, - cleanPatternTokens.length, - 5, - ); - - let prefixMatches = 0; - for (let i = 0; i < maxPrefixCheck; i += 1) { - const tokenSim = cmp.compare(windowTokens[i], cleanPatternTokens[i]); - if (tokenSim < 0.8) break; - prefixMatches += 1; - } - - if (prefixMatches > 0) { - const prefixRatio = prefixMatches / maxPrefixCheck; - boostedRating = adjustedRating * (1 + prefixRatio * 0.25); - } - - if ( - boostedRating > bestRating - || ( - Math.abs(boostedRating - bestRating) < 1e-3 - && lengthDiff < bestLengthDiff - ) - ) { - bestRating = boostedRating; - bestLengthDiff = lengthDiff; - bestStart = start; - bestEnd = start + offset; - } - } - } - - return { - bestStart, - bestEnd, - rating: bestRating, - lengthDiff: bestLengthDiff, - }; -} diff --git a/src/lib/client/pdf-highlight-worker.ts b/src/lib/client/pdf-highlight-worker.ts index 6b396af..fa3fbd6 100644 --- a/src/lib/client/pdf-highlight-worker.ts +++ b/src/lib/client/pdf-highlight-worker.ts @@ -1,6 +1,6 @@ /// -import { findBestHighlightTokenMatch } from './pdf-highlight-match'; +import { findBestHighlightTokenMatch } from './highlight-token-alignment'; interface TokenMatchRequest { id: string; @@ -43,7 +43,10 @@ self.onmessage = (event: MessageEvent) => { const response: TokenMatchResponse = { id, type: 'tokenMatchResult', - ...result, + bestStart: result.start, + bestEnd: result.end, + rating: result.rating, + lengthDiff: result.lengthDiff, }; (self as unknown as DedicatedWorkerGlobalScope).postMessage(response); diff --git a/src/lib/client/pdf.ts b/src/lib/client/pdf.ts index b81234d..1293957 100644 --- a/src/lib/client/pdf.ts +++ b/src/lib/client/pdf.ts @@ -3,11 +3,12 @@ import { TextLayer } from 'pdfjs-dist'; import "core-js/proposals/promise-with-resolvers"; import type { TTSSentenceAlignment } from '@/types/tts'; import type { ParsedPdfDocument, ParsedPdfPage } from '@/types/parsed-pdf'; -import { CmpStr } from 'cmpstr'; import type { TTSSegmentLocator } from '@/types/client'; -import { normalizeUnicodeToken, segmentWords } from '@/lib/shared/language'; - -const cmp = CmpStr.create().setMetric('dice').setFlags('itw'); +import { segmentWords } from '@/lib/shared/language'; +import { + buildAlignmentTokenRanges, + type HighlightTokenRange, +} from '@/lib/client/highlight-token-alignment'; // Worker coordination for offloading highlight token matching interface HighlightTokenMatchRequest { @@ -152,7 +153,7 @@ let lastSpanNodes: HTMLElement[] = []; let lastTokens: PDFToken[] = []; let lastSentenceTokenWindow: { start: number; end: number } | null = null; let lastSentencePattern: string | null = null; -let lastSentenceWordToTokenMap: number[] | null = null; +let lastSentenceWordToTokenRangeMap: Array | null = null; function getOrCreateHighlightLayer(span: HTMLElement): { layer: HTMLElement; @@ -181,9 +182,6 @@ function getOrCreateHighlightLayer(span: HTMLElement): { return { layer, pageElement, pageRect: pageElement.getBoundingClientRect() }; } -const normalizeWordForMatch = (text: string): string => - normalizeUnicodeToken(text); - // Highlighting functions let highlightPatternSeq = 0; @@ -421,7 +419,7 @@ export function highlightPattern( const cleanPattern = pattern.trim().replace(/\s+/g, ' '); if (!cleanPattern) return; lastSentencePattern = cleanPattern; - lastSentenceWordToTokenMap = null; + lastSentenceWordToTokenRangeMap = null; lastSentenceTokenWindow = null; const parsedDocument = options?.parsedDocument ?? null; const locator = options?.locator ?? null; @@ -672,158 +670,61 @@ export function highlightWordIndex( const end = lastSentenceTokenWindow.end; if (end < start) return; - // Lazily build or refresh the mapping from alignment word - // indices to PDF token indices for this sentence window. + // Lazily build or refresh the shared mapping from alignment words to PDF + // token ranges for this sentence window. if ( - !lastSentenceWordToTokenMap || - lastSentenceWordToTokenMap.length !== words.length + !lastSentenceWordToTokenRangeMap || + lastSentenceWordToTokenRangeMap.length !== words.length ) { - const pdfFiltered: { tokenIndex: number; norm: string }[] = []; - for (let i = start; i <= end; i++) { - const norm = normalizeWordForMatch(lastTokens[i].text); - if (!norm) continue; - pdfFiltered.push({ tokenIndex: i, norm }); - } - - const ttsFiltered: { wordIndex: number; norm: string }[] = []; - for (let i = 0; i < words.length; i++) { - const norm = normalizeWordForMatch(words[i].text); - if (!norm) continue; - ttsFiltered.push({ wordIndex: i, norm }); - } - - const wordToToken = new Array(words.length).fill(-1); - - const m = pdfFiltered.length; - const n = ttsFiltered.length; - - if (m && n) { - const dp: number[][] = Array.from({ length: m + 1 }, () => - new Array(n + 1).fill(Number.POSITIVE_INFINITY) - ); - const bt: number[][] = Array.from({ length: m + 1 }, () => - new Array(n + 1).fill(0) - ); // 0=diag,1=up,2=left - - dp[0][0] = 0; - const GAP_COST = 0.7; - - for (let i = 0; i <= m; i++) { - for (let j = 0; j <= n; j++) { - if (i > 0 && j > 0) { - const a = pdfFiltered[i - 1].norm; - const b = ttsFiltered[j - 1].norm; - const sim = cmp.compare(a, b); - const subCost = 1 - sim; - const cand = dp[i - 1][j - 1] + subCost; - if (cand < dp[i][j]) { - dp[i][j] = cand; - bt[i][j] = 0; - } - } - if (i > 0) { - const cand = dp[i - 1][j] + GAP_COST; - if (cand < dp[i][j]) { - dp[i][j] = cand; - bt[i][j] = 1; - } - } - if (j > 0) { - const cand = dp[i][j - 1] + GAP_COST; - if (cand < dp[i][j]) { - dp[i][j] = cand; - bt[i][j] = 2; - } - } - } - } - - let i = m; - let j = n; - while (i > 0 || j > 0) { - const move = bt[i][j]; - if (i > 0 && j > 0 && move === 0) { - const pdfIdx = pdfFiltered[i - 1].tokenIndex; - const ttsIdx = ttsFiltered[j - 1].wordIndex; - if (wordToToken[ttsIdx] === -1) { - wordToToken[ttsIdx] = pdfIdx; - } - i -= 1; - j -= 1; - } else if (i > 0 && (move === 1 || j === 0)) { - i -= 1; - } else if (j > 0 && (move === 2 || i === 0)) { - j -= 1; - } else { - break; - } - } - - // Propagate nearest known mapping to fill gaps - let lastSeen = -1; - for (let k = 0; k < wordToToken.length; k++) { - if (wordToToken[k] !== -1) { - lastSeen = wordToToken[k]; - } else if (lastSeen !== -1) { - wordToToken[k] = lastSeen; - } - } - let nextSeen = -1; - for (let k = wordToToken.length - 1; k >= 0; k--) { - if (wordToToken[k] !== -1) { - nextSeen = wordToToken[k]; - } else if (nextSeen !== -1) { - wordToToken[k] = nextSeen; - } - } - } - - lastSentenceWordToTokenMap = wordToToken; + const relativeRanges = buildAlignmentTokenRanges( + words, + lastTokens.slice(start, end + 1).map((token) => token.text), + { fillGaps: true }, + ); + lastSentenceWordToTokenRangeMap = relativeRanges.map((range) => ( + range ? { start: range.start + start, end: range.end + start } : null + )); } - const mappedIndex = - lastSentenceWordToTokenMap && wordIndex < lastSentenceWordToTokenMap.length - ? lastSentenceWordToTokenMap[wordIndex] - : -1; + const tokenRange = lastSentenceWordToTokenRangeMap[wordIndex]; + if (!tokenRange) return; - if (mappedIndex === -1) return; + for (let tokenIndex = tokenRange.start; tokenIndex <= tokenRange.end; tokenIndex += 1) { + const token = lastTokens[tokenIndex]; + const span = lastSpanNodes[token.spanIndex]; + if (!span) continue; - const chosenTokenIndex = mappedIndex; + const node = token.textNode; + if (!node || node.nodeType !== Node.TEXT_NODE) continue; - const token = lastTokens[chosenTokenIndex]; - const span = lastSpanNodes[token.spanIndex]; - if (!span) return; + try { + const range = document.createRange(); + range.setStart(node, token.startOffset); + range.setEnd(node, token.endOffset); - const node = token.textNode; - if (!node || node.nodeType !== Node.TEXT_NODE) return; + const highlightTarget = getOrCreateHighlightLayer(span); + if (!highlightTarget) continue; - try { - const range = document.createRange(); - range.setStart(node, token.startOffset); - range.setEnd(node, token.endOffset); + const { layer: highlightLayer, pageRect } = highlightTarget; + const rects = Array.from(range.getClientRects()); - const highlightTarget = getOrCreateHighlightLayer(span); - if (!highlightTarget) return; - - const { layer: highlightLayer, pageRect } = highlightTarget; - const rects = Array.from(range.getClientRects()); - - rects.forEach((rect) => { - const highlight = document.createElement('div'); - highlight.className = 'pdf-word-highlight-overlay'; - highlight.style.position = 'absolute'; - highlight.style.backgroundColor = 'var(--accent)'; - highlight.style.opacity = '0.4'; - highlight.style.pointerEvents = 'none'; - highlight.style.left = `${rect.left - pageRect.left}px`; - highlight.style.top = `${rect.top - pageRect.top}px`; - highlight.style.width = `${rect.width}px`; - highlight.style.height = `${rect.height}px`; - highlight.style.zIndex = '2'; - highlightLayer.appendChild(highlight); - }); - } catch { - // Ignore range errors + rects.forEach((rect) => { + const highlight = document.createElement('div'); + highlight.className = 'pdf-word-highlight-overlay'; + highlight.style.position = 'absolute'; + highlight.style.backgroundColor = 'var(--accent)'; + highlight.style.opacity = '0.4'; + highlight.style.pointerEvents = 'none'; + highlight.style.left = `${rect.left - pageRect.left}px`; + highlight.style.top = `${rect.top - pageRect.top}px`; + highlight.style.width = `${rect.width}px`; + highlight.style.height = `${rect.height}px`; + highlight.style.zIndex = '2'; + highlightLayer.appendChild(highlight); + }); + } catch { + // Ignore range errors + } } } diff --git a/tests/unit/epub-word-highlight.vitest.spec.ts b/tests/unit/epub-word-highlight.vitest.spec.ts index 20f0266..6d0e3cc 100644 --- a/tests/unit/epub-word-highlight.vitest.spec.ts +++ b/tests/unit/epub-word-highlight.vitest.spec.ts @@ -1,7 +1,6 @@ import { describe, expect, test } from 'vitest'; import { - buildMonotonicWordToTokenMap, resolveAlignmentWordSourceRange, tokenizeCanonicalSegment, } from '../../src/lib/client/epub/epub-word-highlight'; @@ -19,15 +18,6 @@ const segment = (text: string, offset = 0): CanonicalTtsSegment => ({ spansSourceBoundary: false, }); -const alignmentWords = (words: string[]): TTSSentenceAlignment['words'] => - words.map((word, index) => ({ - text: word, - startSec: index, - endSec: index + 0.5, - charStart: 0, - charEnd: word.length, - })); - describe('EPUB word highlight mapping', () => { test('tokenizes Japanese and Chinese using locale-aware word boundaries', () => { const japanese = tokenizeCanonicalSegment(segment('これは日本語です。', 5), 'ja'); @@ -78,23 +68,4 @@ describe('EPUB word highlight mapping', () => { ]); }); - test('maps repeated words monotonically instead of jumping to later duplicates', () => { - const tokens = tokenizeCanonicalSegment(segment('the light and the shadow and the light returned')); - const map = buildMonotonicWordToTokenMap( - alignmentWords(['the', 'light', 'and', 'the', 'shadow', 'and', 'the', 'light', 'returned']), - tokens, - ); - - expect(map).toEqual([0, 1, 2, 3, 4, 5, 6, 7, 8]); - }); - - test('leaves unmatched alignment words unhighlighted instead of borrowing a neighbor', () => { - const tokens = tokenizeCanonicalSegment(segment('alpha beta gamma')); - const map = buildMonotonicWordToTokenMap( - alignmentWords(['alpha', 'missing', 'gamma']), - tokens, - ); - - expect(map).toEqual([0, -1, 2]); - }); }); diff --git a/tests/unit/highlight-token-alignment.vitest.spec.ts b/tests/unit/highlight-token-alignment.vitest.spec.ts new file mode 100644 index 0000000..7ee3212 --- /dev/null +++ b/tests/unit/highlight-token-alignment.vitest.spec.ts @@ -0,0 +1,92 @@ +import { describe, expect, test } from 'vitest'; + +import { + buildAlignmentTokenRanges, + findBestHighlightTokenMatch, +} from '../../src/lib/client/highlight-token-alignment'; +import { segmentWords } from '../../src/lib/shared/language'; +import type { TTSSentenceAlignment } from '../../src/types/tts'; + +const alignmentWords = ( + words: string[], +): TTSSentenceAlignment['words'] => + words.map((word, index) => ({ + text: word, + startSec: index, + endSec: index + 0.5, + charStart: 0, + charEnd: word.length, + })); + +describe('shared viewer highlight token alignment', () => { + test('matches a complete Japanese sentence using locale-aware token count', () => { + const sentence = 'これは日本語です。'; + const patternTokens = segmentWords(sentence, 'ja').map((token) => token.text); + const tokenTexts = ['前文', ...patternTokens, '次文']; + + expect(findBestHighlightTokenMatch(patternTokens, tokenTexts)).toMatchObject({ + start: 1, + end: patternTokens.length, + lengthDiff: 0, + }); + }); + + test('matches spaced Latin text without relying on whitespace tokens', () => { + expect(findBestHighlightTokenMatch( + ['hello', 'world'], + ['before', 'hello', 'world', 'after'], + )).toMatchObject({ + start: 1, + end: 2, + lengthDiff: 0, + }); + }); + + test('maps a Japanese timed chunk across every visible token it spans', () => { + expect(buildAlignmentTokenRanges( + alignmentWords(['これは', '日本語', 'です']), + ['これ', 'は', '日本語', 'です'], + )).toEqual([ + { start: 0, end: 1 }, + { start: 2, end: 2 }, + { start: 3, end: 3 }, + ]); + }); + + test('maps visible chunks back to a larger timed Japanese token', () => { + expect(buildAlignmentTokenRanges( + alignmentWords(['これ', 'は', '日本語', 'です']), + ['これは', '日本語', 'です'], + )).toEqual([ + { start: 0, end: 0 }, + { start: 0, end: 0 }, + { start: 1, end: 1 }, + { start: 2, end: 2 }, + ]); + }); + + test('maps repeated words monotonically', () => { + expect(buildAlignmentTokenRanges( + alignmentWords(['the', 'light', 'and', 'the', 'light']), + ['the', 'light', 'and', 'the', 'light'], + )).toEqual([ + { start: 0, end: 0 }, + { start: 1, end: 1 }, + { start: 2, end: 2 }, + { start: 3, end: 3 }, + { start: 4, end: 4 }, + ]); + }); + + test('can leave unrelated fallback tokens unmapped for strict viewers', () => { + expect(buildAlignmentTokenRanges( + alignmentWords(['alpha', 'missing', 'gamma']), + ['alpha', 'beta', 'gamma'], + { minimumSimilarity: 0.8 }, + )).toEqual([ + { start: 0, end: 0 }, + null, + { start: 2, end: 2 }, + ]); + }); +}); diff --git a/tests/unit/pdf-highlight-match.vitest.spec.ts b/tests/unit/pdf-highlight-match.vitest.spec.ts deleted file mode 100644 index 716fc26..0000000 --- a/tests/unit/pdf-highlight-match.vitest.spec.ts +++ /dev/null @@ -1,29 +0,0 @@ -import { describe, expect, test } from 'vitest'; - -import { findBestHighlightTokenMatch } from '../../src/lib/client/pdf-highlight-match'; -import { segmentWords } from '../../src/lib/shared/language'; - -describe('PDF highlight token matching', () => { - test('matches a complete Japanese sentence using locale-aware token count', () => { - const sentence = 'これは日本語です。'; - const patternTokens = segmentWords(sentence, 'ja').map((token) => token.text); - const tokenTexts = ['前文', ...patternTokens, '次文']; - - expect(findBestHighlightTokenMatch(patternTokens, tokenTexts)).toMatchObject({ - bestStart: 1, - bestEnd: patternTokens.length, - lengthDiff: 0, - }); - }); - - test('matches spaced Latin text without relying on whitespace tokens', () => { - expect(findBestHighlightTokenMatch( - ['hello', 'world'], - ['before', 'hello', 'world', 'after'], - )).toMatchObject({ - bestStart: 1, - bestEnd: 2, - lengthDiff: 0, - }); - }); -});