diff --git a/packages/core/src/offsets.ts b/packages/core/src/offsets.ts new file mode 100644 index 0000000..7c4a114 --- /dev/null +++ b/packages/core/src/offsets.ts @@ -0,0 +1,72 @@ +/** + * Helpers for moving between the WASM splitter's coordinate system and + * JavaScript's. + * + * `split_offsets` and `chunk_offsets` return **UTF-8 byte** offsets, while + * `String.prototype.slice` indexes **UTF-16 code units**. The two only agree + * for ASCII text, so byte offsets must be converted before they touch a string. + */ + +/** + * Number of UTF-8 bytes needed to encode a single code point. + */ +function utf8Length(codePoint: number): number { + if (codePoint < 0x80) return 1; + if (codePoint < 0x800) return 2; + if (codePoint < 0x10000) return 3; + return 4; +} + +/** + * Convert UTF-8 byte offsets into UTF-16 code-unit offsets for the given text. + * + * Walks the string once, so the cost is linear in the length of the text + * regardless of how many offsets are supplied. ASCII-only text is detected up + * front and returned untouched, which keeps the common case free. + * + * @param text - The text the offsets refer to + * @param byteOffsets - Array of [start, end] UTF-8 byte offset pairs + * @returns Array of [start, end] UTF-16 code-unit offset pairs + */ +export function toCharOffsets( + text: string, + byteOffsets: Array<[number, number]> +): Array<[number, number]> { + if (byteOffsets.length === 0) return []; + + // Byte and code-unit offsets coincide for ASCII, which is the common case. + if (!/[^\x00-\x7F]/.test(text)) return byteOffsets; + + const targets = new Set(); + for (const [start, end] of byteOffsets) { + targets.add(start); + targets.add(end); + } + + const sorted = [...targets].sort((a, b) => a - b); + const charByByte = new Map(); + + let bytePos = 0; + let unitPos = 0; + let cursor = 0; + + for (const char of text) { + while (cursor < sorted.length && sorted[cursor] <= bytePos) { + charByByte.set(sorted[cursor], unitPos); + cursor++; + } + bytePos += utf8Length(char.codePointAt(0)!); + unitPos += char.length; + } + + // Any remaining target is at or past the end of the text. + while (cursor < sorted.length) { + charByByte.set(sorted[cursor], unitPos); + cursor++; + } + + return byteOffsets.map(([start, end]) => [ + charByByte.get(start)!, + charByByte.get(end)!, + ]); +} diff --git a/packages/core/src/recursive.ts b/packages/core/src/recursive.ts index 6f36338..2c8131a 100644 --- a/packages/core/src/recursive.ts +++ b/packages/core/src/recursive.ts @@ -1,4 +1,5 @@ import { init as initChunk, split_offsets, merge_splits } from '@chonkiejs/chunk'; +import { toCharOffsets } from './offsets'; import { Tokenizer } from '@/tokenizer'; import { Chunk, RecursiveRules, RecursiveLevel, IncludeDelim } from '@/types'; @@ -134,11 +135,11 @@ export class RecursiveChunker { private splitText(text: string, level: RecursiveLevel): string[] { // Whitespace splitting - use WASM split with space delimiter if (level.whitespace) { - const offsets = split_offsets(text, { + const offsets = toCharOffsets(text, split_offsets(text, { delimiters: ' ', includeDelim: 'none', minChars: 0 - }); + })); return offsets.map(([start, end]) => text.slice(start, end)); } @@ -159,11 +160,11 @@ export class RecursiveChunker { level.includeDelim === 'prev' ? 'prev' : level.includeDelim === 'next' ? 'next' : 'none'; - const offsets = split_offsets(text, { + const offsets = toCharOffsets(text, split_offsets(text, { delimiters: delims, includeDelim, minChars: this.minCharactersPerChunk - }); + })); return offsets.map(([start, end]) => text.slice(start, end)); } diff --git a/packages/core/src/semantic.ts b/packages/core/src/semantic.ts index a664af7..594f093 100644 --- a/packages/core/src/semantic.ts +++ b/packages/core/src/semantic.ts @@ -7,6 +7,7 @@ */ import { init as initChunk, split_offsets } from '@chonkiejs/chunk'; +import { toCharOffsets } from './offsets'; import { Tokenizer } from '@/tokenizer'; import { Chunk } from '@/types'; @@ -369,11 +370,11 @@ export class SemanticChunker { const raw = this.delimiters.join(''); const delimChars = [...new Set(raw)].filter(c => c !== ' ').join(''); - const offsets = split_offsets(text, { + const offsets = toCharOffsets(text, split_offsets(text, { delimiters: delimChars, includeDelim: this.includeDelim === 'none' ? 'none' : this.includeDelim, minChars: this.minCharactersPerSentence, - }); + })); if (offsets.length === 0) return []; diff --git a/packages/core/src/sentence.ts b/packages/core/src/sentence.ts index 5cabcd4..157d758 100644 --- a/packages/core/src/sentence.ts +++ b/packages/core/src/sentence.ts @@ -3,6 +3,7 @@ */ import { split_offsets, merge_splits } from '@chonkiejs/chunk'; +import { toCharOffsets } from './offsets'; import { initWasm } from '@/wasm'; import { Tokenizer } from '@/tokenizer'; import { Chunk, IncludeDelim } from '@/types'; @@ -141,11 +142,11 @@ export class SentenceChunker { // All single-byte delimiters: use WASM split_offsets const delimStr = this.delim.join(''); - return split_offsets(text, { + return toCharOffsets(text, split_offsets(text, { delimiters: delimStr, includeDelim: this.includeDelim, minChars: this.minCharactersPerSentence, - }); + })); } /** diff --git a/packages/core/tests/unicode-offsets.test.ts b/packages/core/tests/unicode-offsets.test.ts new file mode 100644 index 0000000..8b2ae43 --- /dev/null +++ b/packages/core/tests/unicode-offsets.test.ts @@ -0,0 +1,147 @@ +import { RecursiveChunker, SentenceChunker } from '../src'; +import { toCharOffsets } from '../src/offsets'; + +// The WASM splitter reports UTF-8 byte offsets while String.prototype.slice +// indexes UTF-16 code units. The two agree only for ASCII, so these tests pair +// each non-ASCII input with an ASCII input of identical code-unit length: +// chunking must not care which one it is given. +// 'é' / '–' -> 1 code unit, 2-3 UTF-8 bytes (ASCII twin: 'e' / '-') +// '🔗' -> 2 code units, 4 UTF-8 bytes (ASCII twin: 'ab') +const PARITY_CASES = { + accented: { wide: 'café '.repeat(400), ascii: 'cafe '.repeat(400) }, + dashed: { wide: 'a – b '.repeat(400), ascii: 'a - b '.repeat(400) }, + emoji: { wide: '🔗 link '.repeat(300), ascii: 'ab link '.repeat(300) }, +}; + +describe('Unicode offset handling', () => { + describe('toCharOffsets', () => { + it('should leave ASCII offsets untouched', () => { + expect(toCharOffsets('ab cd', [[0, 2], [3, 5]])).toEqual([[0, 2], [3, 5]]); + }); + + it('should convert byte offsets for two-byte characters', () => { + expect(toCharOffsets('é cd', [[0, 2], [3, 5]])).toEqual([[0, 1], [2, 4]]); + }); + + it('should convert byte offsets for astral characters', () => { + expect(toCharOffsets('🔗 cd', [[0, 4], [5, 7]])).toEqual([[0, 2], [3, 5]]); + }); + + it('should handle an empty offset list', () => { + expect(toCharOffsets('🔗', [])).toEqual([]); + }); + }); + + describe('RecursiveChunker', () => { + // A long line without sentence delimiters forces the recursion down to the + // whitespace level, which is where offsets are applied to the string. + for (const [name, { wide, ascii }] of Object.entries(PARITY_CASES)) { + for (const chunkSize of [64, 256, 1024]) { + it(`should chunk ${name} text like its ASCII twin (chunkSize ${chunkSize})`, async () => { + const chunker = await RecursiveChunker.create({ + chunkSize, + minCharactersPerChunk: 16, + }); + + const wideChunks = await chunker.chunk(wide); + const asciiChunks = await chunker.chunk(ascii); + + expect(wideChunks.map(chunk => chunk.text.length)).toEqual( + asciiChunks.map(chunk => chunk.text.length) + ); + expect(wideChunks.map(chunk => chunk.startIndex)).toEqual( + asciiChunks.map(chunk => chunk.startIndex) + ); + }); + } + } + + it('should not drop characters from the middle of the text', async () => { + const chunker = await RecursiveChunker.create({ + chunkSize: 64, + minCharactersPerChunk: 16, + }); + const text = 'café '.repeat(400); + + const chunks = await chunker.chunk(text); + + // The whitespace level drops the final delimiter for ASCII text too, so + // compare against the trimmed source rather than the raw input. + expect(chunks.map(chunk => chunk.text).join('')).toBe(text.trimEnd()); + }); + + it('should never split a surrogate pair', async () => { + const chunker = await RecursiveChunker.create({ + chunkSize: 64, + minCharactersPerChunk: 16, + }); + + const chunks = await chunker.chunk(PARITY_CASES.emoji.wide); + + for (const chunk of chunks) { + expect(chunk.text.isWellFormed()).toBe(true); + } + }); + + it('should report offsets that index back into the source text', async () => { + const chunker = await RecursiveChunker.create({ + chunkSize: 64, + minCharactersPerChunk: 16, + }); + const text = 'café '.repeat(400); + + const chunks = await chunker.chunk(text); + + for (const chunk of chunks) { + expect(text.slice(chunk.startIndex, chunk.endIndex)).toBe(chunk.text); + } + }); + }); + + describe('SentenceChunker', () => { + // Single-character delimiters take the WASM path; the multi-character + // defaults are split in JS and were never affected. + const delim = ['.', '!', '?']; + const text = 'Le café était très bon. Un peu cher. Mais bon quand même. Voilà.'; + + it('should not run offsets past the end of the text', async () => { + const chunker = await SentenceChunker.create({ + chunkSize: 32, + minCharactersPerSentence: 4, + delim, + }); + + const chunks = await chunker.chunk(text); + + expect(chunks[chunks.length - 1].endIndex).toBe(text.length); + }); + + it('should split sentences at delimiters rather than mid-word', async () => { + const chunker = await SentenceChunker.create({ + chunkSize: 32, + minCharactersPerSentence: 4, + delim, + }); + + const chunks = await chunker.chunk(text); + + for (const chunk of chunks) { + expect(chunk.text.trimEnd().endsWith('.')).toBe(true); + } + }); + + it('should never split a surrogate pair', async () => { + const chunker = await SentenceChunker.create({ + chunkSize: 64, + minCharactersPerSentence: 8, + delim, + }); + + const chunks = await chunker.chunk('Voir le lien 🔗 ici. '.repeat(60)); + + for (const chunk of chunks) { + expect(chunk.text.isWellFormed()).toBe(true); + } + }); + }); +});