diff --git a/packages/transformers/src/pipelines/token-classification.js b/packages/transformers/src/pipelines/token-classification.js index 70914606e..fdbbb1107 100644 --- a/packages/transformers/src/pipelines/token-classification.js +++ b/packages/transformers/src/pipelines/token-classification.js @@ -118,9 +118,10 @@ export class TokenClassificationPipeline } const isBatched = Array.isArray(texts); + const textList = isBatched ? texts : [texts]; // Run tokenization - const model_inputs = this.tokenizer(isBatched ? texts : [texts], { + const model_inputs = this.tokenizer(textList, { padding: true, truncation: true, }); @@ -136,26 +137,40 @@ export class TokenClassificationPipeline for (let i = 0; i < logits.dims[0]; ++i) { const ids = model_inputs.input_ids[i].tolist(); const batch = logits[i]; + const text = textList[i]; const tokens = []; + let charOffset = 0; for (let j = 0; j < batch.dims[0]; ++j) { const tokenData = batch[j]; const topScoreIndex = max(tokenData.data)[1]; const entity = id2label ? id2label[topScoreIndex] : `LABEL_${topScoreIndex}`; - if (ignore_labels.includes(entity)) continue; // TODO add option to keep special tokens? const word = this.tokenizer.decode([ids[j]], { skip_special_tokens: true }); if (word === '') continue; // Was a special token. + // Locate this token's character span in the original text by + // scanning forward from where the previous token ended. + const idx = text.indexOf(word, charOffset); + let start, end; + if (idx !== -1) { + start = idx; + end = idx + word.length; + charOffset = end; + } + + if (ignore_labels.includes(entity)) continue; + const scores = softmax(tokenData.data); tokens.push({ entity, score: scores[topScoreIndex], index: j, word, - // TODO: Add support for start and end + start, + end, }); } @@ -218,10 +233,13 @@ function groupEntities(tokens, ids, tokenizer) { scoreSum += tokens[i].score; groupIds.push(ids[tokens[i].index]); } + const charStart = tokens[start].start; + const charEnd = tokens[end - 1].end; return { entity_group: tag, score: scoreSum / (end - start), word: tokenizer.decode(groupIds, { skip_special_tokens: true }), + ...(charStart !== undefined ? { start: charStart, end: charEnd } : {}), }; }); } diff --git a/packages/transformers/src/tokenization_utils.js b/packages/transformers/src/tokenization_utils.js index 29de6f186..83a362b68 100644 --- a/packages/transformers/src/tokenization_utils.js +++ b/packages/transformers/src/tokenization_utils.js @@ -101,6 +101,42 @@ const SPECIAL_TOKEN_ATTRIBUTES = [ * @param {string} side Which side to pad the array. * @private */ +/** + * Compute character-level [start, end) offsets for each token by scanning + * forward through the original text. Tokens that cannot be found (e.g. + * special tokens like [CLS]/[SEP], or subwords after normalization) get + * [0, 0], matching the Python tokenizers convention. + * + * The scan is tried case-sensitively first, then case-insensitively, to + * handle uncased tokenizers that lowercase the input before tokenizing. + * + * @param {string[]} tokens The token strings produced by the tokenizer. + * @param {string} text The original input text. + * @returns {[number, number][]} + */ +function computeOffsets(tokens, text) { + /** @type {[number, number][]} */ + const offsets = []; + const textLower = text.toLowerCase(); + let pos = 0; + for (const token of tokens) { + if (token === '') { + offsets.push([0, 0]); + continue; + } + // Try exact match first, then case-insensitive for uncased tokenizers. + let idx = text.indexOf(token, pos); + if (idx === -1) idx = textLower.indexOf(token.toLowerCase(), pos); + if (idx === -1) { + offsets.push([0, 0]); + } else { + offsets.push([idx, idx + token.length]); + pos = idx + token.length; + } + } + return offsets; +} + function padHelper(item, length, value_fn, side) { for (const key of Object.keys(item)) { const diff = length - item[key].length; @@ -197,6 +233,7 @@ function getSpecialTokens(tokenizer) { * @property {number|null} [max_length=null] Maximum length of the returned list and optionally padding length. * @property {TReturnTensor} [return_tensor=true] Whether to return the results as Tensors or arrays. * @property {boolean|null} [return_token_type_ids=null] Whether to return the token type ids. + * @property {boolean} [return_offsets_mapping=false] Whether to return character-level [start, end) offsets for each token. */ /** @@ -359,7 +396,7 @@ export class PreTrainedTokenizer text, options = {}, ) { - const { text_pair = null, add_special_tokens = true, padding = false, return_token_type_ids = null } = options; + const { text_pair = null, add_special_tokens = true, padding = false, return_token_type_ids = null, return_offsets_mapping = false } = options; let { truncation = null, max_length = null } = options; const return_tensor = /** @type {TReturnTensor} */ (options.return_tensor ?? true); // Different to HF @@ -380,10 +417,10 @@ export class PreTrainedTokenizer } encodedTokens = text.map((t, i) => - this._encode_plus(t, { text_pair: text_pair[i], add_special_tokens, return_token_type_ids }), + this._encode_plus(t, { text_pair: text_pair[i], add_special_tokens, return_token_type_ids, return_offsets_mapping }), ); } else { - encodedTokens = text.map((x) => this._encode_plus(x, { add_special_tokens, return_token_type_ids })); + encodedTokens = text.map((x) => this._encode_plus(x, { add_special_tokens, return_token_type_ids, return_offsets_mapping })); } } else { if (text === null || text === undefined) { @@ -397,7 +434,7 @@ export class PreTrainedTokenizer } // For single input, we just wrap in an array, and then unwrap later. - encodedTokens = [this._encode_plus(text, { text_pair, add_special_tokens, return_token_type_ids })]; + encodedTokens = [this._encode_plus(text, { text_pair, add_special_tokens, return_token_type_ids, return_offsets_mapping })]; } // At this point, `encodedTokens` is batched, of shape [batch_size, tokens]. // However, array may be jagged. So, we may need pad to max_length. @@ -444,7 +481,7 @@ export class PreTrainedTokenizer padHelper( encodedTokens[i], max_length, - (key) => (key === 'input_ids' ? this.pad_token_id : 0), + (key) => (key === 'input_ids' ? this.pad_token_id : key === 'offset_mapping' ? [0, 0] : 0), this.padding_side, ); } @@ -454,6 +491,12 @@ export class PreTrainedTokenizer const result = {}; + // offset_mapping is a number[][] — it cannot be tensorized. + // Extract it before the tensor loop and re-attach as a plain array. + const offsetMappings = return_offsets_mapping + ? encodedTokens.map((x) => { const v = x.offset_mapping; delete x.offset_mapping; return v; }) + : null; + if (return_tensor) { if (!(padding && truncation)) { // Not, guaranteed that all items have same length, so @@ -502,7 +545,11 @@ export class PreTrainedTokenizer } } - return /** @type {BatchEncoding>} */ (result); + if (offsetMappings) { + result.offset_mapping = isBatched ? offsetMappings : offsetMappings[0]; + } + + return /** @type {BatchEncoding>} */ (/** @type {unknown} */ (result)); } /** @@ -524,11 +571,12 @@ export class PreTrainedTokenizer * @param {string|null} [options.text_pair=null] The optional second text to encode. * @param {boolean} [options.add_special_tokens=true] Whether or not to add the special tokens associated with the corresponding model. * @param {boolean|null} [options.return_token_type_ids=null] Whether to return token_type_ids. - * @returns {{input_ids: number[], attention_mask: number[], token_type_ids?: number[]}} An object containing the encoded text. + * @param {boolean} [options.return_offsets_mapping=false] Whether to return character-level [start, end) offsets for each token. + * @returns {{input_ids: number[], attention_mask: number[], token_type_ids?: number[], offset_mapping?: [number, number][]}} An object containing the encoded text. * @private */ - _encode_plus(text, { text_pair = null, add_special_tokens = true, return_token_type_ids = null } = {}) { - const { ids, attention_mask, token_type_ids } = this._tokenizer.encode(text, { + _encode_plus(text, { text_pair = null, add_special_tokens = true, return_token_type_ids = null, return_offsets_mapping = false } = {}) { + const { ids, attention_mask, token_type_ids, tokens } = this._tokenizer.encode(text, { text_pair, add_special_tokens, return_token_type_ids: return_token_type_ids ?? this.return_token_type_ids, @@ -537,6 +585,7 @@ export class PreTrainedTokenizer input_ids: ids, attention_mask, ...(token_type_ids ? { token_type_ids } : {}), + ...(return_offsets_mapping ? { offset_mapping: computeOffsets(tokens, text) } : {}), }; } diff --git a/packages/transformers/tests/pipelines/test_pipelines_token_classification.js b/packages/transformers/tests/pipelines/test_pipelines_token_classification.js index 9dc0b8b16..d59a35b73 100644 --- a/packages/transformers/tests/pipelines/test_pipelines_token_classification.js +++ b/packages/transformers/tests/pipelines/test_pipelines_token_classification.js @@ -30,21 +30,24 @@ export default () => { score: 0.5292708, index: 1, word: "1", - // 'start': 0, 'end': 1 + start: 0, + end: 1, }, { entity: "LABEL_0", score: 0.5353687, index: 2, word: "2", - // 'start': 2, 'end': 3 + start: 2, + end: 3, }, { entity: "LABEL_1", score: 0.51381934, index: 3, word: "3", - // 'start': 4, 'end': 5 + start: 4, + end: 5, }, ]; expect(output).toBeCloseToNested(target, 5); @@ -61,7 +64,8 @@ export default () => { score: 0.51381934, index: 3, word: "3", - // 'start': 4, 'end': 5 + start: 4, + end: 5, }, ]; expect(output).toBeCloseToNested(target, 5); @@ -82,21 +86,24 @@ export default () => { score: 0.5292708, index: 1, word: "1", - // 'start': 0, 'end': 1 + start: 0, + end: 1, }, { entity: "LABEL_0", score: 0.5353687, index: 2, word: "2", - // 'start': 2, 'end': 3 + start: 2, + end: 3, }, { entity: "LABEL_1", score: 0.51381934, index: 3, word: "3", - // 'start': 4, 'end': 5 + start: 4, + end: 5, }, ], [ @@ -105,14 +112,16 @@ export default () => { score: 0.5432807, index: 1, word: "4", - // 'start': 0, 'end': 1 + start: 0, + end: 1, }, { entity: "LABEL_1", score: 0.5007693, index: 2, word: "5", - // 'start': 2, 'end': 3 + start: 2, + end: 3, }, ], ]; @@ -131,7 +140,8 @@ export default () => { score: 0.51381934, index: 3, word: "3", - // 'start': 4, 'end': 5 + start: 4, + end: 5, }, ], [ @@ -140,7 +150,8 @@ export default () => { score: 0.5007693, index: 2, word: "5", - // 'start': 2, 'end': 3 + start: 2, + end: 3, }, ], ]; @@ -179,10 +190,10 @@ export default () => { const output = await pipe(inputs, { aggregation_strategy: "simple" }); const target = [ [ - { entity_group: "PER", score: 0.5292708, word: "1" }, - { entity_group: "PER", score: 0.524594, word: "2 3" }, + { entity_group: "PER", score: 0.5292708, word: "1", start: 0, end: 1 }, + { entity_group: "PER", score: 0.524594, word: "2 3", start: 2, end: 5 }, ], - [{ entity_group: "PER", score: 0.52202505, word: "4 5" }], + [{ entity_group: "PER", score: 0.52202505, word: "4 5", start: 0, end: 3 }], ]; expect(output).toBeCloseToNested(target, 5); }, @@ -212,7 +223,10 @@ export default () => { "aggregation_strategy='simple' drops O-labeled tokens", async () => { const output = await pipe(inputs, { aggregation_strategy: "simple" }); - const target = [[{ entity_group: "PER", score: 0.51381934, word: "3" }], [{ entity_group: "PER", score: 0.5007693, word: "5" }]]; + const target = [ + [{ entity_group: "PER", score: 0.51381934, word: "3", start: 4, end: 5 }], + [{ entity_group: "PER", score: 0.5007693, word: "5", start: 2, end: 3 }], + ]; expect(output).toBeCloseToNested(target, 5); }, MAX_TEST_EXECUTION_TIME, @@ -224,12 +238,12 @@ export default () => { const output = await pipe(inputs, { aggregation_strategy: "simple", ignore_labels: [] }); const target = [ [ - { entity_group: "O", score: 0.5323198, word: "1 2" }, - { entity_group: "PER", score: 0.51381934, word: "3" }, + { entity_group: "O", score: 0.5323198, word: "1 2", start: 0, end: 3 }, + { entity_group: "PER", score: 0.51381934, word: "3", start: 4, end: 5 }, ], [ - { entity_group: "O", score: 0.5432808, word: "4" }, - { entity_group: "PER", score: 0.5007693, word: "5" }, + { entity_group: "O", score: 0.5432808, word: "4", start: 0, end: 1 }, + { entity_group: "PER", score: 0.5007693, word: "5", start: 2, end: 3 }, ], ]; expect(output).toBeCloseToNested(target, 5); @@ -266,12 +280,12 @@ export default () => { // Labels for `4 5`: [E-PER, B-PER]. const target = [ [ - { entity_group: "PER", score: 0.5323198, word: "1 2" }, - { entity_group: "PER", score: 0.51381934, word: "3" }, + { entity_group: "PER", score: 0.5323198, word: "1 2", start: 0, end: 3 }, + { entity_group: "PER", score: 0.51381934, word: "3", start: 4, end: 5 }, ], [ - { entity_group: "PER", score: 0.5432808, word: "4" }, - { entity_group: "PER", score: 0.5007693, word: "5" }, + { entity_group: "PER", score: 0.5432808, word: "4", start: 0, end: 1 }, + { entity_group: "PER", score: 0.5007693, word: "5", start: 2, end: 3 }, ], ]; expect(output).toBeCloseToNested(target, 5); @@ -305,7 +319,10 @@ export default () => { "aggregation_strategy='simple' folds I-* / E-* into one group per terminator", async () => { const output = await pipe(inputs, { aggregation_strategy: "simple" }); - const target = [[{ entity_group: "PER", score: 0.52614963, word: "1 2 3" }], [{ entity_group: "PER", score: 0.522025, word: "4 5" }]]; + const target = [ + [{ entity_group: "PER", score: 0.52614963, word: "1 2 3", start: 0, end: 5 }], + [{ entity_group: "PER", score: 0.522025, word: "4 5", start: 0, end: 3 }], + ]; expect(output).toBeCloseToNested(target, 5); }, MAX_TEST_EXECUTION_TIME, @@ -337,13 +354,13 @@ export default () => { const output = await pipe(inputs, { aggregation_strategy: "simple" }); const target = [ [ - { entity_group: "PER", score: 0.5292708, word: "1" }, - { entity_group: "PER", score: 0.5353687, word: "2" }, - { entity_group: "PER", score: 0.51381934, word: "3" }, + { entity_group: "PER", score: 0.5292708, word: "1", start: 0, end: 1 }, + { entity_group: "PER", score: 0.5353687, word: "2", start: 2, end: 3 }, + { entity_group: "PER", score: 0.51381934, word: "3", start: 4, end: 5 }, ], [ - { entity_group: "PER", score: 0.5432808, word: "4" }, - { entity_group: "PER", score: 0.5007693, word: "5" }, + { entity_group: "PER", score: 0.5432808, word: "4", start: 0, end: 1 }, + { entity_group: "PER", score: 0.5007693, word: "5", start: 2, end: 3 }, ], ]; expect(output).toBeCloseToNested(target, 5); diff --git a/packages/transformers/tests/tokenizers.test.js b/packages/transformers/tests/tokenizers.test.js index 1e7977ba9..cadc10e3d 100644 --- a/packages/transformers/tests/tokenizers.test.js +++ b/packages/transformers/tests/tokenizers.test.js @@ -457,6 +457,59 @@ describe("Token type ids", () => { ); }); +describe("Offset mapping", () => { + it( + "single string — returns [start, end) for each token", + async () => { + const tokenizer = await AutoTokenizer.from_pretrained("Xenova/bert-base-uncased"); + + const output = tokenizer("Hello world", { + return_tensor: false, + return_offsets_mapping: true, + }); + + // bert-base-uncased adds [CLS] and [SEP] as special tokens (empty string → [0,0]) + expect(output.offset_mapping).toEqual([ + [0, 0], // [CLS] + [0, 5], // "hello" + [6, 11], // "world" + [0, 0], // [SEP] + ]); + }, + MAX_TEST_EXECUTION_TIME, + ); + + it( + "batched strings — returns an array of offset arrays", + async () => { + const tokenizer = await AutoTokenizer.from_pretrained("Xenova/bert-base-uncased"); + + const output = tokenizer(["Hi", "a b"], { + padding: true, + truncation: true, + return_tensor: false, + return_offsets_mapping: true, + }); + + expect(output.offset_mapping).toEqual([ + [[0, 0], [0, 2], [0, 0], [0, 0]], // "Hi" padded to length 4 + [[0, 0], [0, 1], [2, 3], [0, 0]], // "a b" + ]); + }, + MAX_TEST_EXECUTION_TIME, + ); + + it( + "offset_mapping is absent when return_offsets_mapping is not set", + async () => { + const tokenizer = await AutoTokenizer.from_pretrained("Xenova/bert-base-uncased"); + const output = tokenizer("Hello", { return_tensor: false }); + expect(output.offset_mapping).toBeUndefined(); + }, + MAX_TEST_EXECUTION_TIME, + ); +}); + describe("Edge cases", () => { it( "should not crash when encoding a very long string",