add pre tokenization option to token classification pipeline

Author: riccardopinosio

backends/pipeline.go

@@ -125,6 +125,7 @@ type TokenizedInput struct {
 	AttentionMask     []uint32
 	SpecialTokensMask []uint32
 	Offsets           [][2]uint
+	WordIDs           []int
 	MaxAttentionIndex int
 }
 

hugot_test.go

@@ -724,6 +724,19 @@ func tokenClassificationPipeline(t *testing.T, session *Session) {
 	pipelineNone, err3 := NewPipeline(session, configNone)
 	checkT(t, err3)
 
+	// Split-words enabled pipeline
+	configSplit := TokenClassificationConfig{
+		ModelPath: modelPath,
+		Name:      "testPipelineSplitWords",
+		Options: []TokenClassificationOption{
+			pipelines.WithSimpleAggregation(),
+			pipelines.WithIgnoreLabels([]string{"O"}),
+			pipelines.WithSplitWords(),
+		},
+	}
+	pipelineSplit, errSplit := NewPipeline(session, configSplit)
+	checkT(t, errSplit)
+
 	var expectedResults map[int]pipelines.TokenClassificationOutput
 	err4 := json.Unmarshal(embedded.TokenExpectedByte, &expectedResults)
 	checkT(t, err4)
@@ -769,6 +782,65 @@ func tokenClassificationPipeline(t *testing.T, session *Session) {
 			}
 		})
 	}
+
+	// Expect same entities as the simple aggregation for the equivalent sentence for split words
+	t.Run("Split words aggregation", func(t *testing.T) {
+		words := [][]string{{"My", "name", "is", "Wolfgang", "and", "I", "live", "in", "Berlin", "."}}
+		batchResult, err := pipelineSplit.RunWords(words)
+		checkT(t, err)
+		printTokenEntities(batchResult)
+		expected := expectedResults[0]
+		for i, predictedEntities := range batchResult.Entities {
+			assert.Equal(t, len(expected.Entities[i]), len(predictedEntities))
+			for j, entity := range predictedEntities {
+				expectedEntity := expected.Entities[i][j]
+				assert.Equal(t, expectedEntity.Entity, entity.Entity)
+				assert.Equal(t, expectedEntity.Word, entity.Word)
+			}
+		}
+	})
+
+	t.Run("Split words yields different offsets vs double-space input", func(t *testing.T) {
+		// Non-split input with double space changes raw offsets.
+		nonSplit := []string{"New  York is great."}
+		splitWords := [][]string{{"New", "York", "is", "great", "."}}
+
+		resNonSplit, errNon := pipelineSimple.RunPipeline(nonSplit)
+		checkT(t, errNon)
+		resSplit, errSplitRun := pipelineSplit.RunWords(splitWords)
+		checkT(t, errSplitRun)
+
+		// Compare first sequence: entity words may match, but offsets should differ due to space normalization.
+		if len(resNonSplit.Entities) > 0 && len(resSplit.Entities) > 0 && len(resNonSplit.Entities[0]) > 0 && len(resSplit.Entities[0]) > 0 {
+			eNon := resNonSplit.Entities[0][0]
+			eSplit := resSplit.Entities[0][0]
+			assert.NotEqual(t, fmt.Sprintf("%d-%d", eNon.Start, eNon.End), fmt.Sprintf("%d-%d", eSplit.Start, eSplit.End))
+		}
+	})
+
+	// Pre-tokenization splitting on 'X': expect different entity results from the non-split case
+	t.Run("Split on X detects entity", func(t *testing.T) {
+		nonSplit := []string{"XBerlinXXisXbeautiful."}
+		split := [][]string{{"Berlin is", "beautiful", "."}}
+
+		resNonSplit, errNS := pipelineSimple.RunPipeline(nonSplit)
+		checkT(t, errNS)
+		resSplit, errSW := pipelineSplit.RunWords(split)
+		checkT(t, errSW)
+
+		gotA := resNonSplit.Entities[0]
+		gotB := resSplit.Entities[0]
+		// Expect split-words to detect 'Berlin' as an entity, while non-split should not because of the confusing X characters.
+		hasBerlin := func(es []pipelines.Entity) bool {
+			for _, e := range es {
+				if strings.EqualFold(e.Word, "Berlin") {
+					return true
+				}
+			}
+			return false
+		}
+		assert.True(t, !hasBerlin(gotA) || len(gotA) < len(gotB), "expected split-words to surface 'Berlin' or increase entity count")
+	})
 }
 
 func tokenClassificationPipelineValidation(t *testing.T, session *Session) {

pipelines/tokenClassification.go

@@ -21,6 +21,7 @@ type TokenClassificationPipeline struct {
 	IDLabelMap          map[int]string
 	AggregationStrategy string
 	IgnoreLabels        []string
+	SplitWords          bool
 }
 type Entity struct {
 	Entity    string
@@ -98,6 +99,14 @@ func WithIgnoreLabels(ignoreLabels []string) backends.PipelineOption[*TokenClass
 	}
 }
 
+// WithSplitWords enables word-level alignment like Hugging Face's is_split_into_words.
+func WithSplitWords() backends.PipelineOption[*TokenClassificationPipeline] {
+	return func(pipeline *TokenClassificationPipeline) error {
+		pipeline.SplitWords = true
+		return nil
+	}
+}
+
 // NewTokenClassificationPipeline Initializes a feature extraction pipeline.
 func NewTokenClassificationPipeline(config backends.PipelineConfig[*TokenClassificationPipeline], s *options.Options, model *backends.Model) (*TokenClassificationPipeline, error) {
 	defaultPipeline, err := backends.NewBasePipeline(config, s, model)
@@ -179,6 +188,9 @@ func (p *TokenClassificationPipeline) Validate() error {
 
 // Preprocess tokenizes the input strings.
 func (p *TokenClassificationPipeline) Preprocess(batch *backends.PipelineBatch, inputs []string) error {
+	if p.SplitWords {
+		return fmt.Errorf("split-words enabled: use RunWords/PreprocessWords for [][]string inputs")
+	}
 	start := time.Now()
 	backends.TokenizeInputs(batch, p.Model.Tokenizer, inputs)
 	atomic.AddUint64(&p.Model.Tokenizer.TokenizerTimings.NumCalls, 1)
@@ -187,6 +199,79 @@ func (p *TokenClassificationPipeline) Preprocess(batch *backends.PipelineBatch,
 	return err
 }
 
+// PreprocessWords tokenizes pre-split words and maps tokens to word IDs via offsets.
+func (p *TokenClassificationPipeline) PreprocessWords(batch *backends.PipelineBatch, inputs [][]string) error {
+	start := time.Now()
+	// Join words with single spaces to simulate pretokenized behavior
+	joined := make([]string, len(inputs))
+	wordBoundaries := make([][][2]uint, len(inputs))
+	// local helper to convert non-negative int to uint safely
+	toUintNonNeg := func(i int) uint {
+		if i < 0 {
+			return 0
+		}
+		return uint(i)
+	}
+	for i, words := range inputs {
+		joined[i] = strings.Join(words, " ")
+		// compute boundaries in joined string
+		var boundaries [][2]uint
+		pos := 0
+		for wIdx, w := range words {
+			startPos := pos
+			endPos := pos + len(w)
+			// clamp to non-negative and convert safely to uint
+			// ensure non-negative before converting to uint
+			if startPos < 0 {
+				startPos = 0
+			}
+			if endPos < 0 {
+				endPos = 0
+			}
+			boundaries = append(boundaries, [2]uint{toUintNonNeg(startPos), toUintNonNeg(endPos)})
+			// add one space after every word except last
+			if wIdx < len(words)-1 {
+				pos = endPos + 1
+			} else {
+				pos = endPos
+			}
+		}
+		wordBoundaries[i] = boundaries
+	}
+	backends.TokenizeInputs(batch, p.Model.Tokenizer, joined)
+	atomic.AddUint64(&p.Model.Tokenizer.TokenizerTimings.NumCalls, 1)
+	atomic.AddUint64(&p.Model.Tokenizer.TokenizerTimings.TotalNS, safeconv.DurationToU64(time.Since(start)))
+
+	// Map token offsets to word indices
+	for i := range batch.Input {
+		input := batch.Input[i]
+		boundaries := wordBoundaries[i]
+		wordIDs := make([]int, len(input.Offsets))
+		for t := range input.Offsets {
+			if input.SpecialTokensMask[t] > 0 {
+				wordIDs[t] = -1
+				continue
+			}
+			tokStart := input.Offsets[t][0]
+			tokEnd := input.Offsets[t][1]
+			id := -1
+			for w := range boundaries {
+				b := boundaries[w]
+				// assign if token lies within the word span
+				if tokStart >= b[0] && tokEnd <= b[1] {
+					id = w
+					break
+				}
+			}
+			wordIDs[t] = id
+		}
+		batch.Input[i].WordIDs = wordIDs
+		// also set raw to joined string for offsets consistency
+		batch.Input[i].Raw = joined[i]
+	}
+	return backends.CreateInputTensors(batch, p.Model, p.Runtime)
+}
+
 // Forward performs the forward inference of the pipeline.
 func (p *TokenClassificationPipeline) Forward(batch *backends.PipelineBatch) error {
 	start := time.Now()
@@ -257,6 +342,7 @@ func (p *TokenClassificationPipeline) GatherPreEntities(input backends.Tokenized
 		endInd := input.Offsets[j][1]
 		wordRef := sentence[startInd:endInd]
 		isSubword := len(word) != len(wordRef)
+		// In split-words mode, grouping will use offsets between tokens rather than IsSubword.
 		// TODO: check for unknown token here, it's in the config and can be loaded and compared with the token
 		// in that case set the subword as in the python code
 		preEntities = append(preEntities, Entity{
@@ -355,15 +441,29 @@ func (p *TokenClassificationPipeline) aggregateWords(entities []Entity) ([]Entit
 	for _, entity := range entities {
 		if len(wordGroup) == 0 {
 			wordGroup = []Entity{entity}
-		} else if entity.IsSubword {
-			wordGroup = append(wordGroup, entity)
-		} else {
+			continue
+		}
+		// Default behavior: group by IsSubword boundaries
+		groupBreak := !entity.IsSubword
+		if p.SplitWords {
+			// In split-words mode, we group by contiguous tokens of the same word boundary since we pretokenized.
+			// Since preEntities don’t carry word IDs directly, simulate grouping by contiguous offsets:
+			// break group if there is a gap between previous End and current Start (space) or heuristic non-subword.
+			// TODO: eventually we should export word IDs from the tokenizer to avoid this heuristic but the rust tokenizer bindings don't expose this yet
+			// and we also use other tokenizers in go backend.
+			prev := wordGroup[len(wordGroup)-1]
+			// if there is a gap in offsets consider it a new word
+			groupBreak = entity.Start > prev.End
+		}
+		if groupBreak {
 			aggregated, err := p.aggregateWord(wordGroup)
 			if err != nil {
 				return nil, err
 			}
 			wordEntities = append(wordEntities, aggregated)
 			wordGroup = []Entity{entity}
+		} else {
+			wordGroup = append(wordGroup, entity)
 		}
 	}
 	if len(wordGroup) > 0 {
@@ -516,3 +616,26 @@ func (p *TokenClassificationPipeline) RunPipeline(inputs []string) (*TokenClassi
 	runErrors = append(runErrors, postErr)
 	return result, errors.Join(runErrors...)
 }
+
+// RunWords runs the pipeline for pre-split word inputs.
+// Each input is a slice of words representing a pretokenized sentence.
+// This is particularly useful when the user wants to control tokenization because of special tokens,
+// hashtags, or other domain-specific tokenization needs.
+func (p *TokenClassificationPipeline) RunWords(inputs [][]string) (*TokenClassificationOutput, error) {
+	var runErrors []error
+	batch := backends.NewBatch(len(inputs))
+	defer func(*backends.PipelineBatch) {
+		runErrors = append(runErrors, batch.Destroy())
+	}(batch)
+	runErrors = append(runErrors, p.PreprocessWords(batch, inputs))
+	if e := errors.Join(runErrors...); e != nil {
+		return nil, e
+	}
+	runErrors = append(runErrors, p.Forward(batch))
+	if e := errors.Join(runErrors...); e != nil {
+		return nil, e
+	}
+	result, postErr := p.Postprocess(batch)
+	runErrors = append(runErrors, postErr)
+	return result, errors.Join(runErrors...)
+}