Spaces:
Running
Running
| /** | |
| * Pulpie WebUI - Block Chunker | |
| * | |
| * Implements token sequence packing matching Pulpie's exact chunker algorithm. | |
| * Groups blocks into sequences bounded by max_tokens (default 8,192). | |
| * | |
| * Sequence structure: | |
| * [BOS] block_0 [<|sep|>] block_1 [<|sep|>] ... block_N [<|sep|>] [EOS] | |
| */ | |
| import { TOKEN_IDS } from './tokenizer.js'; | |
| export const DEFAULT_MAX_TOKENS = 1024; | |
| /** | |
| * @typedef {Object} Chunk | |
| * @property {number[]} chunkIds - Array of token IDs including BOS, EOS, and SEP tokens | |
| * @property {number[]} blockIndices - Indices of the blocks included in this chunk | |
| */ | |
| /** | |
| * Pack block tokens into chunks conforming to EuroBERT's max_tokens limit. | |
| * | |
| * @param {number[][]} blockTokenIds - Array of token ID lists for each block | |
| * @param {Object} [options] | |
| * @param {number} [options.maxTokens=8192] - Maximum sequence length per forward pass | |
| * @param {number} [options.sepTokenId] - Token ID for <|sep|> (default 128256) | |
| * @param {number} [options.bosTokenId] - Token ID for <|begin_of_text|> (default 128000) | |
| * @param {number} [options.eosTokenId] - Token ID for <|end_of_text|> (default 128001) | |
| * @returns {Chunk[]} | |
| */ | |
| export function packChunks( | |
| blockTokenIds, | |
| { | |
| maxTokens = DEFAULT_MAX_TOKENS, | |
| sepTokenId = TOKEN_IDS.SEP, | |
| bosTokenId = TOKEN_IDS.BOS, | |
| eosTokenId = TOKEN_IDS.EOS, | |
| } = {} | |
| ) { | |
| if (!blockTokenIds || blockTokenIds.length === 0) { | |
| return []; | |
| } | |
| const chunks = []; | |
| let currentTokens = []; | |
| let currentBlockIndices = []; | |
| const overhead = (bosTokenId !== null ? 1 : 0) + (eosTokenId !== null ? 1 : 0); | |
| const budget = maxTokens - overhead; | |
| let currentLen = 0; | |
| for (let blockIdx = 0; blockIdx < blockTokenIds.length; blockIdx++) { | |
| const blockToks = blockTokenIds[blockIdx]; | |
| const blockCost = blockToks.length + 1; // +1 for SEP token | |
| if (currentLen + blockCost > budget) { | |
| if (currentTokens.length > 0) { | |
| const chunkIds = []; | |
| if (bosTokenId !== null) chunkIds.push(bosTokenId); | |
| chunkIds.push(...currentTokens); | |
| if (eosTokenId !== null) chunkIds.push(eosTokenId); | |
| chunks.push({ | |
| chunkIds, | |
| blockIndices: [...currentBlockIndices], | |
| }); | |
| } | |
| currentTokens = []; | |
| currentBlockIndices = []; | |
| currentLen = 0; | |
| // Single block exceeds total budget on its own | |
| if (blockCost > budget) { | |
| const truncated = blockToks.slice(0, budget - 1); | |
| currentTokens.push(...truncated, sepTokenId); | |
| currentBlockIndices.push(blockIdx); | |
| currentLen = truncated.length + 1; | |
| continue; | |
| } | |
| } | |
| currentTokens.push(...blockToks, sepTokenId); | |
| currentBlockIndices.push(blockIdx); | |
| currentLen += blockCost; | |
| } | |
| if (currentTokens.length > 0) { | |
| const chunkIds = []; | |
| if (bosTokenId !== null) chunkIds.push(bosTokenId); | |
| chunkIds.push(...currentTokens); | |
| if (eosTokenId !== null) chunkIds.push(eosTokenId); | |
| chunks.push({ | |
| chunkIds, | |
| blockIndices: [...currentBlockIndices], | |
| }); | |
| } | |
| return chunks; | |
| } | |