Added new HuggingFaceTokenizer to token_utils and updated TokenizerParameters to include tokenizer_type and access_token as additional metadata to store. (#229)

Author: gpolovets1

benchmarks/benchmark_serving.py

@@ -208,6 +208,7 @@ def get_tokenizer(
     model_id: str,
     tokenizer_name: str,
     use_hf_tokenizer: bool,
+    access_token: str | None = None,
 ) -> Any:
   """Return a tokenizer or a tokenizer placholder."""
   if tokenizer_name == "test":
@@ -218,7 +219,7 @@ def get_tokenizer(
     # follow up instructions below to set up access token
     # https://huggingface.co/docs/transformers.js/en/guides/private
     print(f"Using HuggingFace tokenizer: {tokenizer_name}")
-    return AutoTokenizer.from_pretrained(tokenizer_name)
+    return AutoTokenizer.from_pretrained(tokenizer_name, token=access_token)
   elif model_id == "llama-3":
     # Llama 3 uses a tiktoken tokenizer.
     print(f"Using llama-3 tokenizer: {tokenizer_name}")
@@ -386,11 +387,9 @@ def load_math500_dataset(dataset_path: str) -> list[tuple[Any, Any]]:
 
 
 def tokenize_dataset(
-    dataset: list[tuple[Any, Any, Any]],
-    tokenizer: Any,
+    dataset: list[tuple[Any, Any, Any]], tokenizer: Any, use_chat_template: bool
 ) -> list[tuple[str, Any, str, int, int, int]]:
   tokenized_dataset = []
-
   for prompt, output, idx in dataset:
     if isinstance(output, tuple):
       output_len = len(tokenizer.encode(output[0]))
@@ -399,7 +398,12 @@ def tokenize_dataset(
       output_len = len(tokenizer.encode(output))
       output_tokens = output
 
-    prompt_tokens = tokenizer.encode(prompt)
+    if use_chat_template:
+      prompt_tokens = tokenizer.apply_chat_template(
+          [{"role": "user", "content": prompt}], add_generation_prompt=True
+      )
+    else:
+      prompt_tokens = tokenizer.encode(prompt)
 
     tokenized_data = (
         prompt,
@@ -470,6 +474,7 @@ def filter_dataset(
 def sample_requests(
     dataset: list[tuple[Any, Any]],
     tokenizer: Any,
+    use_chat_template: bool,
     num_requests: int,
     dataset_type: str,
     max_output_length: int = 0,
@@ -508,7 +513,9 @@ def sample_requests(
     sampled_data = dataset[i] + (dataset_indices[i],)
     sampled_dataset.append(sampled_data)
 
-  tokenized_dataset = tokenize_dataset(sampled_dataset, tokenizer)
+  tokenized_dataset = tokenize_dataset(
+      sampled_dataset, tokenizer, use_chat_template
+  )
 
   input_requests = filter_dataset(
       tokenized_dataset,
@@ -636,6 +643,7 @@ async def grpc_async_request(
 async def send_request(
     api_url: str,
     tokenizer: Any,
+    use_chat_template: bool,
     input_request: InputRequest,
     prefill_quota: AsyncCounter,
     active_req_quota: AsyncCounter,
@@ -645,7 +653,13 @@ async def send_request(
 ) -> RequestFuncOutput:
   """Send the request to JetStream server."""
   # Tokenize on client side following MLPerf standard.
-  token_ids = tokenizer.encode(input_request.prompt)
+  if use_chat_template:
+    token_ids = tokenizer.apply_chat_template(
+        [{"role": "user", "content": input_request.prompt}],
+        add_generation_prompt=True,
+    )
+  else:
+    token_ids = tokenizer.encode(input_request.prompt)
 
   # Send the request
   request = jetstream_pb2.DecodeRequest(
@@ -691,6 +705,7 @@ async def send_request(
 async def benchmark(
     api_url: str,
     tokenizer: Any,
+    use_chat_template: bool,
     input_requests: list[InputRequest],
     request_rate: float,
     disable_tqdm: bool,
@@ -734,6 +749,7 @@ async def benchmark(
             send_request(
                 api_url=api_url,
                 tokenizer=tokenizer,
+                use_chat_template=use_chat_template,
                 input_request=request,
                 prefill_quota=prefill_quota,
                 active_req_quota=active_req_quota,
@@ -892,6 +908,23 @@ def parse_args() -> argparse.Namespace:
           " to True, and provide name of the tokenizer in the tokenizer flag."
       ),
   )
+  parser.add_argument(
+      "--hf-access-token",
+      type=str,
+      default="",
+      help=(
+          "Access token used to load a tokenizer from an API (i.e. HuggingFace)"
+      ),
+  )
+  parser.add_argument(
+      "--use-chat-template",
+      type=str2bool,
+      default=False,
+      help=(
+          "Whether the tokenizer should be applying a chat template "
+          "(used for instruction-tuned models)."
+      ),
+  )
   parser.add_argument(
       "--num-prompts",
       type=int,
@@ -1051,13 +1084,16 @@ def main(args: argparse.Namespace):
   model_id = args.model
   tokenizer_id = args.tokenizer
   use_hf_tokenizer = args.use_hf_tokenizer
+  hf_access_token = args.hf_access_token
+  use_chat_template = args.use_chat_template
 
   prefill_quota = AsyncCounter(init_value=3)
   active_req_quota = AsyncCounter(init_value=450)
 
   api_url = f"{args.server}:{args.port}"
-
-  tokenizer = get_tokenizer(model_id, tokenizer_id, use_hf_tokenizer)
+  tokenizer = get_tokenizer(
+      model_id, tokenizer_id, use_hf_tokenizer, hf_access_token
+  )
   if tokenizer == "test" or args.dataset == "test":
     input_requests = mock_requests(
         args.total_mock_requests
@@ -1094,6 +1130,7 @@ def main(args: argparse.Namespace):
     input_requests = sample_requests(
         dataset=dataset,
         tokenizer=tokenizer,
+        use_chat_template=use_chat_template,
         num_requests=args.num_prompts,
         dataset_type=args.dataset,
         max_output_length=args.max_output_length,
@@ -1116,6 +1153,7 @@ def main(args: argparse.Namespace):
         benchmark(
             api_url=api_url,
             tokenizer=tokenizer,
+            use_chat_template=use_chat_template,
             input_requests=warmup_requests,
             request_rate=args.request_rate,
             disable_tqdm=args.disable_tqdm,
@@ -1134,6 +1172,7 @@ def main(args: argparse.Namespace):
       benchmark(
           api_url=api_url,
           tokenizer=tokenizer,
+          use_chat_template=use_chat_template,
           input_requests=input_requests,
           request_rate=args.request_rate,
           disable_tqdm=args.disable_tqdm,

benchmarks/tests/test_benchmark_serving.py

@@ -27,6 +27,7 @@ class TestBenchmarkServing(unittest.IsolatedAsyncioTestCase):
   async def test_benchmark(self):
     api_url = "test_url"
     tokenizer = mock.MagicMock()
+    use_chat_template = False
     tokenizer.encode = mock.MagicMock(return_value=[1, 2, 3])
     tokenizer.decode = mock.MagicMock(return_value="test_decode")
     input_requests = [
@@ -78,6 +79,7 @@ def mock_orchestrator_factory(*args, **kwargs):
       metrics, outputs = await benchmark_serving.benchmark(
           api_url,
           tokenizer,
+          use_chat_template,
           input_requests,
           request_rate,
           disable_tqdm,

jetstream/engine/mock_engine.py

@@ -450,7 +450,11 @@ def get_prefix_destination_sharding(self) -> Any:
 
   def get_tokenizer(self) -> tokenizer_pb2.TokenizerParameters:
     """Return a protobuf of tokenizer info, callable from Py or C++."""
-    return tokenizer_pb2.TokenizerParameters(path="test", extra_ids=0)
+    return tokenizer_pb2.TokenizerParameters(
+        path="test",
+        tokenizer_type=tokenizer_pb2.TokenizerType.sentencepiece,
+        extra_ids=0,
+    )
 
   def init_decode_state(self) -> DecodeState:
     """Initialises any state which a generation step transforms."""

jetstream/engine/token_utils.py

@@ -18,6 +18,7 @@
 import logging
 from typing import Any, Iterable, List, Optional, Tuple, Union
 
+from transformers import AutoTokenizer
 import jax
 import jax.numpy as jnp
 import numpy as np
@@ -200,7 +201,7 @@ def pad_tokens(
     tokens: Tokens.
     bos_id: Bos ID.
     pad_id: Pad ID.
-    is_bos: Add a beginning of sequence token if this is ture.
+    is_bos: Add a beginning of sequence token if this is true.
     prefill_lengths: Buckets to pad the sequence to for static compilation.
     max_prefill_length: Maximum bucket to use.
     jax_padding: convert to JAX padded tokens if True.
@@ -506,3 +507,103 @@ def eos_id(self) -> int:
   def bos_id(self) -> int:
     """ID of the BOS token."""
     return self.tokenizer.bos_id
+
+
+class HuggingFaceTokenizer(tokenizer_api.Tokenizer):
+  """Tokenizer to convert strings to token ids and vice-versa."""
+
+  def __init__(self, metadata: tokenizer_pb2.TokenizerParameters):
+    self.tokenizer = AutoTokenizer.from_pretrained(
+        metadata.path, token=metadata.access_token
+    )
+    self.metadata = metadata
+
+  def encode(
+      self, s: str, **kwargs
+  ) -> Tuple[Union[jax.Array, np.ndarray], int]:
+    """Tokenize a string.
+    Args:
+        s: String to tokenize.
+        **kwargs: Additional keyword arguments
+    Returns:
+        tokens: Tokenized into integers.
+        true_length: Actual length of the non-padded sequence
+          if padding is used.
+    """
+    is_bos = kwargs.pop("is_bos", True)
+    prefill_lengths = kwargs.pop("prefill_lengths", None)
+    max_prefill_length = kwargs.pop("max_prefill_length", None)
+    jax_padding = kwargs.pop("jax_padding", True)
+    if getattr(self.metadata, "use_chat_template", False):
+      tokens = self.tokenizer.apply_chat_template(
+          [{"role": "user", "content": s}],
+          add_generation_prompt=True,
+          return_tensors="np",
+      ).squeeze()
+      if is_bos:
+        logging.warning(
+            "Overriding is_bos to False because use_chat_template "
+            "is set to True."
+        )
+      is_bos = False
+    else:
+      tokens = self.tokenizer.encode(
+          s, add_special_tokens=False, return_tensors="np"
+      ).squeeze()
+
+    tokens, true_length = pad_tokens(
+        tokens,
+        self.bos_id,
+        self.pad_id,
+        is_bos=is_bos,
+        prefill_lengths=prefill_lengths,
+        max_prefill_length=max_prefill_length,
+        jax_padding=jax_padding,
+    )
+    return tokens, true_length
+
+  def decode(self, token_ids: list[int]) -> str:
+    """Processess input token ids to generate a string.
+    Args:
+      token_ids: List of token ids.
+    Returns:
+      str: String generated from the token ids.
+    """
+    return self.tokenizer.decode(token_ids, skip_special_tokens=True)
+
+  @property
+  def pad_id(self) -> Union[None, int]:
+    """ID of the pad token."""
+    if getattr(self.tokenizer, "pad_token_id", None):
+      return self.tokenizer.pad_token_id
+    elif getattr(self.tokenizer, "pad_token", None):
+      try:
+        return self.tokenizer.encode(self.tokenizer.pad_token)[0]
+      except ValueError as _:
+        raise ValueError(
+            "Tokenizer does not contain a special" " pad token!"
+        ) from None
+
+  @property
+  def eos_id(self) -> Union[None, int]:
+    if getattr(self.tokenizer, "eos_token_id", None):
+      return self.tokenizer.eos_token_id
+    elif getattr(self.tokenizer, "eos_token", None):
+      try:
+        return self.tokenizer.encode(self.tokenizer.eos_token)[0]
+      except ValueError as _:
+        raise ValueError(
+            "Tokenizer does not contain a special " "eos token!"
+        ) from None
+
+  @property
+  def bos_id(self) -> Union[None, int]:
+    if getattr(self.tokenizer, "bos_token_id", None):
+      return self.tokenizer.bos_token_id
+    elif getattr(self.tokenizer, "bos_token", None):
+      try:
+        return self.tokenizer.encode(self.tokenizer.bos_token)[0]
+      except ValueError as _:
+        raise ValueError(
+            "Tokenizer does not contain a special " "bos token!"
+        ) from None

jetstream/engine/tokenizer.proto

@@ -20,5 +20,14 @@ option java_multiple_files = true;
 
 message TokenizerParameters {
   string path = 1;
-  int32 extra_ids = 2;
+  TokenizerType tokenizer_type = 2;
+  string access_token = 3;
+  bool use_chat_template = 4;
+  int32 extra_ids = 5;
+}
+
+enum TokenizerType {
+  tiktoken = 0;
+  sentencepiece = 1;
+  huggingface = 2;
 }