remove model name dependency from pool

Signed-off-by: Sage Ahrac <[email protected]>

Author: sagiahrac

pkg/tokenization/pool.go

@@ -205,7 +205,7 @@ func (pool *Pool) processTask(task Task) error {
 		}
 	}
 
-	tokenIDs, overlapRatio := pool.indexer.FindLongestContainedTokens(task.Prompt, task.ModelName)
+	tokenIDs, overlapRatio := pool.indexer.FindLongestContainedTokens(task.Prompt)
 
 	// if the overlap ratio is low, get the full tokenization
 	if overlapRatio < pool.minPrefixOverlapRatio {
@@ -216,7 +216,7 @@ func (pool *Pool) processTask(task Task) error {
 		}
 
 		// update the indexer with the new tokenization
-		if e := pool.indexer.AddTokenization(task.ModelName, task.Prompt, tokens, offsets); e != nil {
+		if e := pool.indexer.AddTokenization(task.Prompt, tokens, offsets); e != nil {
 			err = fmt.Errorf("tokenization failed for model %s: %w", task.ModelName, e)
 			return err
 		}

pkg/tokenization/pool_test.go

@@ -67,14 +67,14 @@ type MockIndexer struct {
 	mock.Mock
 }
 
-func (m *MockIndexer) AddTokenization(modelName, prompt string, tokens []uint32, offsets []tokenizers.Offset) error {
-	args := m.Called(modelName, prompt, tokens, offsets)
+func (m *MockIndexer) AddTokenization(prompt string, tokens []uint32, offsets []tokenizers.Offset) error {
+	args := m.Called(prompt, tokens, offsets)
 	return args.Error(0)
 }
 
 //nolint:gocritic // unnamedResult: tokens and overlapRatio are self-explanatory from context
-func (m *MockIndexer) FindLongestContainedTokens(prompt, modelName string) ([]uint32, float64) {
-	args := m.Called(prompt, modelName)
+func (m *MockIndexer) FindLongestContainedTokens(prompt string) ([]uint32, float64) {
+	args := m.Called(prompt)
 	tokens := args.Get(0).([]uint32) //nolint:errcheck // unused mock
 	return tokens, 0.0
 }