refactor(nlp): improve text splitting and normalization for TTS blocks

- Rename and enhance text processing functions in nlp.ts for better handling of oversized texts, sentence boundaries, and PDF artifacts
- Update PDFViewer to add layout-aware highlighting with retry logic for sentence and word highlights
- Adjust PDFContext and TTSContext to use new normalized text functions
- Expand unit tests for new splitting behaviors, including long texts and punctuation preferences

src/components/PDFViewer.tsx

@@ -24,6 +24,12 @@ export function PDFViewer({ zoomLevel }: PDFViewerProps) {
   const containerRef = useRef<HTMLDivElement>(null);
   const scaleRef = useRef<number>(1);
   const { containerWidth } = usePDFResize(containerRef);
+  const sentenceHighlightSeqRef = useRef(0);
+  const wordHighlightSeqRef = useRef(0);
+  const sentenceHighlightTimeoutsRef = useRef<ReturnType<typeof setTimeout>[]>([]);
+  const wordHighlightTimeoutsRef = useRef<ReturnType<typeof setTimeout>[]>([]);
+  const lastSentenceLayoutKeyRef = useRef<string>('');
+  const lastWordLayoutKeyRef = useRef<string>('');
 
   // Config context
   const { viewType, pdfHighlightEnabled, pdfWordHighlightEnabled } = useConfig();
@@ -49,6 +55,28 @@ export function PDFViewer({ zoomLevel }: PDFViewerProps) {
     currDocPage,
   } = usePDF();
 
+  const layoutKey = `${zoomLevel}:${containerWidth}:${viewType}:${currDocPage}`;
+
+  const clearSentenceHighlightTimeouts = useCallback(() => {
+    for (const t of sentenceHighlightTimeoutsRef.current) clearTimeout(t);
+    sentenceHighlightTimeoutsRef.current = [];
+  }, []);
+
+  const clearWordHighlightTimeouts = useCallback(() => {
+    for (const t of wordHighlightTimeoutsRef.current) clearTimeout(t);
+    wordHighlightTimeoutsRef.current = [];
+  }, []);
+
+  const scheduleSentenceTimeout = useCallback((fn: () => void, ms: number) => {
+    const t = setTimeout(fn, ms);
+    sentenceHighlightTimeoutsRef.current.push(t);
+  }, []);
+
+  const scheduleWordTimeout = useCallback((fn: () => void, ms: number) => {
+    const t = setTimeout(fn, ms);
+    wordHighlightTimeoutsRef.current.push(t);
+  }, []);
+
   useEffect(() => {
     /*
      * Handles highlighting the current sentence being read by TTS.
@@ -66,17 +94,47 @@ export function PDFViewer({ zoomLevel }: PDFViewerProps) {
       return;
     }
 
-    const highlightTimeout = setTimeout(() => {
-      if (containerRef.current) {
-        highlightPattern(currDocText, currentSentence || '', containerRef as RefObject<HTMLDivElement>);
+    clearSentenceHighlightTimeouts();
+
+    const seq = ++sentenceHighlightSeqRef.current;
+    const isLayoutChange = layoutKey !== lastSentenceLayoutKeyRef.current;
+    lastSentenceLayoutKeyRef.current = layoutKey;
+
+    if (isLayoutChange) {
+      clearHighlights();
+    }
+
+    const tryApply = (attempt: number) => {
+      if (seq !== sentenceHighlightSeqRef.current) return;
+      const container = containerRef.current;
+      if (!container) return;
+      if (!currentSentence) return;
+
+      const spans = container.querySelectorAll('.react-pdf__Page__textContent span');
+      if (!spans.length) {
+        if (attempt < 10) scheduleSentenceTimeout(() => tryApply(attempt + 1), 75);
+        return;
       }
-    }, 200);
+
+      highlightPattern(currDocText, currentSentence, containerRef as RefObject<HTMLDivElement>);
+    };
+
+    scheduleSentenceTimeout(() => tryApply(0), 200);
 
     return () => {
-      clearTimeout(highlightTimeout);
+      clearSentenceHighlightTimeouts();
       clearHighlights();
     };
-  }, [currDocText, currentSentence, highlightPattern, clearHighlights, pdfHighlightEnabled]);
+  }, [
+    currDocText,
+    currentSentence,
+    highlightPattern,
+    clearHighlights,
+    pdfHighlightEnabled,
+    layoutKey,
+    clearSentenceHighlightTimeouts,
+    scheduleSentenceTimeout
+  ]);
 
   // Word-level highlight layered on top of the block highlight
   useEffect(() => {
@@ -102,12 +160,41 @@ export function PDFViewer({ zoomLevel }: PDFViewerProps) {
       return;
     }
 
-    highlightWordIndex(
-      currentSentenceAlignment,
-      currentWordIndex,
-      currentSentence || '',
-      containerRef as RefObject<HTMLDivElement>
-    );
+    clearWordHighlightTimeouts();
+
+    const seq = ++wordHighlightSeqRef.current;
+    const isLayoutChange = layoutKey !== lastWordLayoutKeyRef.current;
+    lastWordLayoutKeyRef.current = layoutKey;
+
+    const tryApplyWord = (attempt: number) => {
+      if (seq !== wordHighlightSeqRef.current) return;
+      const container = containerRef.current;
+      if (!container) return;
+
+      highlightWordIndex(
+        currentSentenceAlignment,
+        currentWordIndex,
+        currentSentence || '',
+        containerRef as RefObject<HTMLDivElement>
+      );
+
+      if (isLayoutChange) {
+        // If we don't see a word overlay yet, the sentence highlight worker may not
+        // have produced `lastSentenceTokenWindow` (or the text layer isn't ready).
+        const overlayCount = container.querySelectorAll('.pdf-word-highlight-overlay').length;
+        if (overlayCount === 0 && attempt < 12) {
+          scheduleWordTimeout(() => tryApplyWord(attempt + 1), 75);
+        }
+      }
+    };
+
+    if (isLayoutChange) {
+      clearWordHighlights();
+      scheduleWordTimeout(() => tryApplyWord(0), 250);
+      return;
+    }
+
+    tryApplyWord(0);
   }, [
     currentWordIndex,
     currentSentence,
@@ -115,7 +202,10 @@ export function PDFViewer({ zoomLevel }: PDFViewerProps) {
     pdfHighlightEnabled,
     pdfWordHighlightEnabled,
     clearWordHighlights,
-    highlightWordIndex
+    highlightWordIndex,
+    layoutKey,
+    clearWordHighlightTimeouts,
+    scheduleWordTimeout
   ]);
 
   // Add page dimensions state

src/contexts/PDFContext.tsx

@@ -30,7 +30,7 @@ import type { PDFDocumentProxy } from 'pdfjs-dist';
 import { getPdfDocument } from '@/lib/dexie';
 import { useTTS } from '@/contexts/TTSContext';
 import { useConfig } from '@/contexts/ConfigContext';
-import { processTextToSentences } from '@/lib/nlp';
+import { normalizeTextForTts } from '@/lib/nlp';
 import { withRetry, getAudiobookStatus, generateTTS, createAudiobookChapter } from '@/lib/client';
 import {
   extractTextFromPDF,
@@ -312,9 +312,7 @@ export function PDFProvider({ children }: { children: ReactNode }) {
         });
         const trimmedText = rawText.trim();
         if (trimmedText) {
-          const processedText = smartSentenceSplitting
-            ? processTextToSentences(trimmedText).join(' ')
-            : trimmedText;
+          const processedText = smartSentenceSplitting ? normalizeTextForTts(trimmedText) : trimmedText;
 
           textPerPage.push(processedText);
           totalLength += processedText.length;
@@ -532,7 +530,7 @@ export function PDFProvider({ children }: { children: ReactNode }) {
       }
 
       const textForTTS = smartSentenceSplitting
-        ? processTextToSentences(trimmedText).join(' ')
+        ? normalizeTextForTts(trimmedText)
         : trimmedText;
 
       // Use logical chapter numbering (index + 1) to match original generation titles

src/contexts/TTSContext.tsx

@@ -37,7 +37,7 @@ import { useAudioContext } from '@/hooks/audio/useAudioContext';
 import { getLastDocumentLocation, setLastDocumentLocation } from '@/lib/dexie';
 import { useBackgroundState } from '@/hooks/audio/useBackgroundState';
 import { withRetry, generateTTS, alignAudio } from '@/lib/client';
-import { preprocessSentenceForAudio, processTextToSentences } from '@/lib/nlp';
+import { preprocessSentenceForAudio, splitTextToTtsBlocks } from '@/lib/nlp';
 import { isKokoroModel } from '@/utils/voice';
 import type {
   TTSLocation,
@@ -379,13 +379,13 @@ export function TTSProvider({ children }: { children: ReactNode }): ReactElement
    * @param {string} text - The text to be processed
    * @returns {Promise<string[]>} Array of processed sentences
    */
-  const processTextToSentencesLocal = useCallback(async (text: string): Promise<string[]> => {
+  const splitTextToTtsBlocksLocal = useCallback(async (text: string): Promise<string[]> => {
     if (text.length < 1) {
       return [];
     }
 
     // Use the shared utility directly instead of making an API call
-    return processTextToSentences(text);
+    return splitTextToTtsBlocks(text);
   }, []);
 
   /**
@@ -579,7 +579,7 @@ export function TTSProvider({ children }: { children: ReactNode }): ReactElement
     abortAudio(true); // Clear pending requests since text is changing
     setIsProcessing(true); // Set processing state before text processing starts
 
-    processTextToSentencesLocal(workingText)
+    splitTextToTtsBlocksLocal(workingText)
       .then(newSentences => {
         if (newSentences.length === 0) {
           console.warn('No sentences found in text');
@@ -658,7 +658,7 @@ export function TTSProvider({ children }: { children: ReactNode }): ReactElement
           duration: 3000,
         });
       });
-  }, [isPlaying, handleBlankSection, abortAudio, processTextToSentencesLocal, pendingRestoreIndex, isEPUB, smartSentenceSplitting]);
+  }, [isPlaying, handleBlankSection, abortAudio, splitTextToTtsBlocksLocal, pendingRestoreIndex, isEPUB, smartSentenceSplitting]);
 
   /**
    * Toggles the playback state between playing and paused

src/lib/nlp.ts

@@ -9,6 +9,113 @@ import nlp from 'compromise';
 
 const MAX_BLOCK_LENGTH = 450;
 
+const splitOversizedText = (text: string, maxLen: number): string[] => {
+  const normalized = text.replace(/\s+/g, ' ').trim();
+  if (!normalized) return [];
+  if (normalized.length <= maxLen) return [normalized];
+
+  const parts: string[] = [];
+  const MAX_OVERFLOW = maxLen; // allow finishing the sentence up to +maxLen chars
+  const CLOSERS = new Set(['"', "'", '”', '’', ')', ']', '}']);
+  const BREAK_CHARS = new Set(['.', '!', '?']);
+  const SOFT_BREAK_CHARS = new Set([';', ':']);
+
+  const findPunctuationCut = (s: string, limit: number): number | null => {
+    for (let i = limit; i >= 0; i--) {
+      const ch = s[i];
+      if (!BREAK_CHARS.has(ch)) continue;
+
+      const prev = i > 0 ? s[i - 1] : '';
+      const next = i + 1 < s.length ? s[i + 1] : '';
+
+      // Avoid splitting inside decimals like 3.14
+      if (ch === '.' && /\d/.test(prev) && /\d/.test(next)) continue;
+
+      let end = i + 1;
+      while (end < s.length && CLOSERS.has(s[end])) end++;
+      const after = end < s.length ? s[end] : '';
+
+      // Allow a boundary at end/whitespace, or common PDF artifact where
+      // the next sentence starts immediately with an uppercase letter.
+      if (!after || /\s/.test(after) || /[A-Z]/.test(after)) return end;
+    }
+    return null;
+  };
+
+  const findForwardPunctuationCut = (
+    s: string,
+    startIndex: number,
+    endIndex: number,
+    chars: Set<string>
+  ): number | null => {
+    const start = Math.max(0, startIndex);
+    const end = Math.min(endIndex, s.length - 1);
+    for (let i = start; i <= end; i++) {
+      const ch = s[i];
+      if (!chars.has(ch)) continue;
+
+      const prev = i > 0 ? s[i - 1] : '';
+      const next = i + 1 < s.length ? s[i + 1] : '';
+
+      if (ch === '.' && /\d/.test(prev) && /\d/.test(next)) continue;
+
+      let cut = i + 1;
+      while (cut < s.length && CLOSERS.has(s[cut])) cut++;
+      const after = cut < s.length ? s[cut] : '';
+
+      if (!after || /\s/.test(after) || /[A-Z]/.test(after)) return cut;
+    }
+    return null;
+  };
+
+  const findSoftPunctuationCut = (s: string, limit: number): number | null => {
+    for (let i = limit; i >= 0; i--) {
+      const ch = s[i];
+      if (!SOFT_BREAK_CHARS.has(ch)) continue;
+
+      let end = i + 1;
+      while (end < s.length && CLOSERS.has(s[end])) end++;
+      const after = end < s.length ? s[end] : '';
+      if (!after || /\s/.test(after) || /[A-Z]/.test(after)) return end;
+    }
+    return null;
+  };
+
+  let remaining = normalized;
+  while (remaining.length > maxLen) {
+    const backwardLimit = Math.min(maxLen, remaining.length - 1);
+    const forwardLimit = Math.min(maxLen + MAX_OVERFLOW, remaining.length - 1);
+
+    let cut =
+      findPunctuationCut(remaining, backwardLimit) ??
+      findForwardPunctuationCut(remaining, maxLen, forwardLimit, BREAK_CHARS) ??
+      findSoftPunctuationCut(remaining, backwardLimit) ??
+      findForwardPunctuationCut(remaining, maxLen, forwardLimit, SOFT_BREAK_CHARS) ??
+      remaining.lastIndexOf(' ', maxLen);
+
+    if (cut === 0 || cut === -1) {
+      // No whitespace or punctuation; hard-cut for extremely long tokens.
+      cut = maxLen;
+    }
+
+    const chunk = remaining.slice(0, cut).trim();
+    if (chunk) parts.push(chunk);
+    remaining = remaining.slice(cut).trim();
+  }
+
+  if (remaining) parts.push(remaining);
+  return parts;
+};
+
+const normalizeSentenceBoundariesForNlp = (text: string): string => {
+  // PDF extraction sometimes yields "...end.Next..." with no whitespace.
+  // Insert a space only when it looks like a sentence boundary (lower/digit before,
+  // uppercase after) to avoid breaking abbreviations like "U.S.A".
+  return text
+    .replace(/([a-z0-9])([.!?])(?=[A-Z])/g, '$1$2 ')
+    .replace(/([a-z0-9][.!?][\"”’)\]])(?=[A-Z])/g, '$1 ');
+};
+
 /**
  * Preprocesses text for audio generation by cleaning up various text artifacts
  * 
@@ -31,14 +138,18 @@ export const preprocessSentenceForAudio = (text: string): string => {
  * @param {string} text - The text to split into sentences
  * @returns {string[]} Array of sentence blocks
  */
-export const splitIntoSentences = (text: string): string[] => {
-  const paragraphs = text.split(/\n+/);
+export const splitTextToTtsBlocks = (text: string): string[] => {
+  // Treat double-newlines as paragraph boundaries; single newlines are usually
+  // just PDF line wrapping and should not force sentence/block boundaries.
+  const paragraphs = text.split(/\n{2,}/);
   const blocks: string[] = [];
 
   for (const paragraph of paragraphs) {
     if (!paragraph.trim()) continue;
 
-    const cleanedText = preprocessSentenceForAudio(paragraph);
+    const cleanedText = normalizeSentenceBoundariesForNlp(
+      preprocessSentenceForAudio(paragraph)
+    );
     const doc = nlp(cleanedText);
     const rawSentences = doc.sentences().out('array') as string[];
     
@@ -49,14 +160,17 @@ export const splitIntoSentences = (text: string): string[] => {
 
     for (const sentence of mergedSentences) {
       const trimmedSentence = sentence.trim();
+      const sentenceParts = splitOversizedText(trimmedSentence, MAX_BLOCK_LENGTH);
 
-      if (currentBlock && (currentBlock.length + trimmedSentence.length + 1) > MAX_BLOCK_LENGTH) {
-        blocks.push(currentBlock.trim());
-        currentBlock = trimmedSentence;
-      } else {
-        currentBlock = currentBlock 
-          ? `${currentBlock} ${trimmedSentence}`
-          : trimmedSentence;
+      for (const sentencePart of sentenceParts) {
+        if (currentBlock && (currentBlock.length + sentencePart.length + 1) > MAX_BLOCK_LENGTH) {
+          blocks.push(currentBlock.trim());
+          currentBlock = sentencePart;
+        } else {
+          currentBlock = currentBlock 
+            ? `${currentBlock} ${sentencePart}`
+            : sentencePart;
+        }
       }
     }
 
@@ -69,29 +183,23 @@ export const splitIntoSentences = (text: string): string[] => {
 };
 
 /**
- * Main sentence processing function that handles both short and long texts
+ * Normalizes text for single-shot TTS generation (e.g., a whole PDF page).
+ * Uses the same logic as `splitTextToTtsBlocks`, but returns a single string.
  * 
  * @param {string} text - The text to process
- * @returns {string[]} Array of processed sentences/blocks
+ * @returns {string} Normalized text
  */
-export const processTextToSentences = (text: string): string[] => {
-  if (!text || text.length < 1) {
-    return [];
-  }
-
-  // Always use the full splitting logic so we consistently respect
-  // sentence boundaries and quoted dialogue, even for shorter texts.
-  return splitIntoSentences(text);
-};
+export const normalizeTextForTts = (text: string): string =>
+  splitTextToTtsBlocks(text).join(' ');
 
 /**
- * Gets raw sentences from text without preprocessing or grouping
- * This is useful for text matching and highlighting
+ * Extracts raw sentence strings from text without preprocessing or block grouping.
+ * Useful for text matching and highlighting.
  * 
  * @param {string} text - The text to extract sentences from
  * @returns {string[]} Array of raw sentences
  */
-export const getRawSentences = (text: string): string[] => {
+export const extractRawSentences = (text: string): string[] => {
   if (!text || text.length < 1) {
     return [];
   }
@@ -111,8 +219,8 @@ export const processTextWithMapping = (text: string): {
   rawSentences: string[];
   sentenceMapping: Array<{ processedIndex: number; rawIndices: number[] }>;
 } => {
-  const rawSentences = getRawSentences(text);
-  const processedSentences = processTextToSentences(text);
+  const rawSentences = extractRawSentences(text);
+  const processedSentences = splitTextToTtsBlocks(text);
   
   // Create a mapping between processed sentences and raw sentences
   const sentenceMapping: Array<{ processedIndex: number; rawIndices: number[] }> = [];
@@ -149,6 +257,7 @@ export const processTextWithMapping = (text: string): {
     sentenceMapping
   };
 }; 
+
 // Helper functions to merge quoted dialogue across sentences
 const countDoubleQuotes = (s: string): number => {
   const matches = s.match(/["“”]/g);
@@ -216,4 +325,4 @@ const mergeQuotedDialogue = (rawSentences: string[]): string[] => {
   }
 
   return result;
-};
+};

tests/unit/nlp.spec.ts

@@ -1,7 +1,7 @@
 import { test, expect } from '@playwright/test';
 import { 
   preprocessSentenceForAudio, 
-  splitIntoSentences, 
+  splitTextToTtsBlocks, 
   processTextWithMapping 
 } from '../../src/lib/nlp';
 
@@ -31,24 +31,24 @@ test.describe('preprocessSentenceForAudio', () => {
   });
 });
 
-test.describe('splitIntoSentences', () => {
+test.describe('splitTextToTtsBlocks', () => {
   test('groups short sentences into single block', () => {
     const input = 'First sentence. Second sentence.';
-    const result = splitIntoSentences(input);
+    const result = splitTextToTtsBlocks(input);
     expect(result).toHaveLength(1);
     expect(result[0]).toBe('First sentence. Second sentence.');
   });
 
   test('merges quoted dialogue (double quotes)', () => {
     const input = 'He said, "This should be one block." and walked away.';
-    const result = splitIntoSentences(input);
+    const result = splitTextToTtsBlocks(input);
     expect(result).toHaveLength(1);
     expect(result[0]).toBe('He said, "This should be one block." and walked away.');
   });
   
   test('merges quoted dialogue (curly quotes)', () => {
       const input = 'She replied, “This also should be merged.” then smiled.';
-      const result = splitIntoSentences(input);
+      const result = splitTextToTtsBlocks(input);
       expect(result).toHaveLength(1);
       expect(result[0]).toBe('She replied, “This also should be merged.” then smiled.');
     });
@@ -64,12 +64,75 @@ test.describe('splitIntoSentences', () => {
     // 5 sentences = 505 chars + 4 spaces = 509 chars (> 450). Should split.
     
     const input = Array(5).fill(sentence).join(' ');
-    const result = splitIntoSentences(input);
+    const result = splitTextToTtsBlocks(input);
     
     expect(result.length).toBeGreaterThan(1);
     // The first block should contain as many as possible
     expect(result[0].length).toBeLessThanOrEqual(450);
   });
+
+  test('splits a single oversized sentence into multiple blocks', () => {
+    // Some PDF pages (e.g. research papers) can extract into one massive "sentence"
+    // with few or no punctuation marks; we still must respect MAX_BLOCK_LENGTH.
+    const input = Array(1200).fill('word').join(' '); // no punctuation
+    const result = splitTextToTtsBlocks(input);
+    expect(result.length).toBeGreaterThan(1);
+    for (const block of result) {
+      expect(block.length).toBeGreaterThan(0);
+      expect(block.length).toBeLessThanOrEqual(450);
+    }
+  });
+
+  test('splits extremely long unbroken tokens', () => {
+    const input = 'A'.repeat(1200); // no spaces, no punctuation
+    const result = splitTextToTtsBlocks(input);
+    expect(result.length).toBeGreaterThan(1);
+    for (const block of result) {
+      expect(block.length).toBeGreaterThan(0);
+      expect(block.length).toBeLessThanOrEqual(450);
+    }
+  });
+
+  test('prefers sentence punctuation when chunking long PDF-like text', () => {
+    const sentences = Array.from({ length: 80 }, (_, i) =>
+      `Sentence ${i} has some filler words to keep the length varying slightly number ${i}.`
+    );
+    // Simulate a common PDF extraction artifact: no whitespace after '.' before the next sentence.
+    const input = sentences.join('');
+    const result = splitTextToTtsBlocks(input);
+    expect(result.length).toBeGreaterThan(1);
+    for (const block of result) {
+      expect(block.length).toBeGreaterThan(0);
+      expect(block.length).toBeLessThanOrEqual(450);
+      expect(block.endsWith('.')).toBe(true);
+    }
+  });
+
+  test('does not treat single newlines as paragraph boundaries', () => {
+    // Many PDFs contain hard-wrapped lines; we should not break blocks/sentences
+    // just because of a newline.
+    const input =
+      'The first line ends with a comma,\n' +
+      'but the sentence continues on the next line and ends here.\n' +
+      'And this is the second sentence.';
+    const result = splitTextToTtsBlocks(input);
+    expect(result).toHaveLength(1);
+    expect(result[0]).toBe(
+      'The first line ends with a comma, but the sentence continues on the next line and ends here. And this is the second sentence.'
+    );
+  });
+
+  test('allows long sentences to reach their ending punctuation', () => {
+    const longSentence =
+      `${'word '.repeat(110)}` + // ~550 chars before period
+      'end.' +
+      ' Next.';
+    const result = splitTextToTtsBlocks(longSentence);
+    // The first block should end at a period, not be cut mid-sentence at a space boundary.
+    expect(result.length).toBeGreaterThan(1);
+    expect(result[0].endsWith('.')).toBe(true);
+    expect(result[0].includes('end.')).toBe(true);
+  });
 });
 
 test.describe('processTextWithMapping', () => {