Adds tokenization task

Signed-off-by: Flavia Beo <[email protected]>

caikit_nlp/modules/text_embedding/embedding.py

@@ -42,6 +42,8 @@
  SentenceSimilarityResult,
  SentenceSimilarityResults,
  SentenceSimilarityScores,
+ Token,
+ TokenizationResults,
 )
 from caikit.interfaces.nlp.tasks import (
  EmbeddingTask,
@@ -50,6 +52,7 @@
  RerankTasks,
  SentenceSimilarityTask,
  SentenceSimilarityTasks,
+ TokenizationTask,
 )
 import alog
 
@@ -120,6 +123,7 @@ class TruncatedTokensTuple(NamedTuple):
  SentenceSimilarityTasks,
  RerankTask,
  RerankTasks,
+ TokenizationTask,
  ],
 )
 class EmbeddingModule(ModuleBase):
@@ -192,6 +196,29 @@ def public_model_info(cls) -> Dict[str, Any]: # pylint: disable=no-self-argumen
  "sentence_embedding_dimension": cls.model.get_sentence_embedding_dimension(),
  }
 
+ @TokenizationTask.taskmethod()
+ def run_tokenizer(
+ self,
+ text: str,
+ ) -> TokenizationResults:
+ """Run tokenization task against the model
+
+ Args:
+ text: str
+ Text to tokenize
+ Returns:
+ TokenizationResults
+ The token count
+ """
+ result = self.model.tokenizer.encode_plus(text, return_offsets_mapping=True)
+
+ mapping = [
+ interv for interv in result.offset_mapping if (interv[1] - interv[0]) > 0
+ ]
+ tokens = [Token(start=i[0], end=i[1], text=text[i[0] : i[1]]) for i in mapping]
+
+ return TokenizationResults(token_count=len(result.input_ids), results=tokens)
+
  @classmethod
  def _get_ipex(cls, ipex_flag):
  """Get IPEX optimization library if enabled and available, else return False