crfm-helm 0.5.1__py3-none-any.whl → 0.5.3__py3-none-any.whl

helm/config/tokenizer_configs.yaml

@@ -17,11 +17,31 @@ tokenizer_configs:
     prefix_token: "<s>"
 
   # AI21
-  - name: ai21/j1
+  - name: ai21/j2-tokenizer
     tokenizer_spec:
-      class_name: "helm.tokenizers.ai21_tokenizer.AI21Tokenizer"
-    end_of_text_token: " "
-    prefix_token: ""
+      class_name: "helm.tokenizers.ai21_tokenizer.AI21LocalTokenizer"
+    end_of_text_token: "<|endoftext|>"
+    prefix_token: "<|startoftext|>"
+  - name: ai21/jamba-tokenizer
+    tokenizer_spec:
+      class_name: "helm.tokenizers.ai21_tokenizer.AI21LocalTokenizer"
+    end_of_text_token: "<|endoftext|>"
+    prefix_token: "<|startoftext|>"
+  - name: ai21/jamba-instruct-tokenizer
+    tokenizer_spec:
+      class_name: "helm.tokenizers.ai21_tokenizer.AI21LocalTokenizer"
+    end_of_text_token: "<|endoftext|>"
+    prefix_token: "<|startoftext|>"
+  - name: ai21/jamba-1.5-mini-tokenizer
+    tokenizer_spec:
+      class_name: "helm.tokenizers.ai21_tokenizer.AI21LocalTokenizer"
+    end_of_text_token: "<|endoftext|>"
+    prefix_token: "<|startoftext|>"
+  - name: ai21/jamba-1.5-large-tokenizer
+    tokenizer_spec:
+      class_name: "helm.tokenizers.ai21_tokenizer.AI21LocalTokenizer"
+    end_of_text_token: "<|endoftext|>"
+    prefix_token: "<|startoftext|>"
 
   # AlephAlpha
   - name: AlephAlpha/luminous-base
@@ -45,6 +65,24 @@ tokenizer_configs:
     end_of_text_token: ""
     prefix_token: ""
 
+  # Alibaba DAMO Academy
+
+  - name: damo/seallm-7b-v2
+    tokenizer_spec:
+      class_name: "helm.tokenizers.huggingface_tokenizer.HuggingFaceTokenizer"
+      args:
+        pretrained_model_name_or_path: SeaLLMs/SeaLLM-7B-v2
+    end_of_text_token: "</s>"
+    prefix_token: "<s>"
+
+  - name: damo/seallm-7b-v2.5
+    tokenizer_spec:
+      class_name: "helm.tokenizers.huggingface_tokenizer.HuggingFaceTokenizer"
+      args:
+        pretrained_model_name_or_path: SeaLLMs/SeaLLM-7B-v2.5
+    end_of_text_token: "<eos>"
+    prefix_token: "<bos>"
+
   # Anthropic
   - name: anthropic/claude
     tokenizer_spec:
@@ -69,7 +107,7 @@ tokenizer_configs:
     tokenizer_spec:
       class_name: "helm.tokenizers.huggingface_tokenizer.HuggingFaceTokenizer"
     end_of_text_token: "</s>"
-    prefix_token: "</s>"
+    prefix_token: "<s>"
   - name: bigscience/T0pp
     tokenizer_spec:
       class_name: "helm.tokenizers.huggingface_tokenizer.HuggingFaceTokenizer"
@@ -77,11 +115,45 @@ tokenizer_configs:
     prefix_token: ""
 
   # Cohere
-  - name: cohere/cohere
+  - name: cohere/command
     tokenizer_spec:
-      class_name: "helm.tokenizers.cohere_tokenizer.CohereTokenizer"
-    end_of_text_token: ""
-    prefix_token: ":"
+      class_name: "helm.tokenizers.cohere_tokenizer.CohereLocalTokenizer"
+    end_of_text_token: "<EOS_TOKEN>"
+    prefix_token: "<BOS_TOKEN>"
+
+  - name: cohere/command-light
+    tokenizer_spec:
+      class_name: "helm.tokenizers.cohere_tokenizer.CohereLocalTokenizer"
+    end_of_text_token: "<EOS_TOKEN>"
+    prefix_token: "<BOS_TOKEN>"
+
+  - name: cohere/command-r
+    tokenizer_spec:
+      class_name: "helm.tokenizers.cohere_tokenizer.CohereLocalTokenizer"
+    end_of_text_token: "<EOS_TOKEN>"
+    prefix_token: "<BOS_TOKEN>"
+
+  - name: cohere/command-r-plus
+    tokenizer_spec:
+      class_name: "helm.tokenizers.cohere_tokenizer.CohereLocalTokenizer"
+    end_of_text_token: "<EOS_TOKEN>"
+    prefix_token: "<BOS_TOKEN>"
+
+  - name: cohere/c4ai-command-r-v01
+    tokenizer_spec:
+      class_name: "helm.tokenizers.huggingface_tokenizer.HuggingFaceTokenizer"
+      args:
+        pretrained_model_name_or_path: CohereForAI/c4ai-command-r-v01
+    end_of_text_token: "<EOS_TOKEN>"
+    prefix_token: "<BOS_TOKEN>"
+
+  - name: cohere/c4ai-command-r-plus
+    tokenizer_spec:
+      class_name: "helm.tokenizers.huggingface_tokenizer.HuggingFaceTokenizer"
+      args:
+        pretrained_model_name_or_path: CohereForAI/c4ai-command-r-plus
+    end_of_text_token: "<EOS_TOKEN>"
+    prefix_token: "<BOS_TOKEN>"
 
   # Databricks
   - name: databricks/dbrx-instruct
@@ -159,6 +231,11 @@ tokenizer_configs:
       class_name: "helm.tokenizers.huggingface_tokenizer.HuggingFaceTokenizer"
     end_of_text_token: "<eos>"
     prefix_token: "<bos>"
+  - name: google/gemma-2-9b
+    tokenizer_spec:
+      class_name: "helm.tokenizers.huggingface_tokenizer.HuggingFaceTokenizer"
+    end_of_text_token: "<eos>"
+    prefix_token: "<bos>"
 
   # Hf-internal-testing
 
@@ -240,6 +317,14 @@ tokenizer_configs:
     prefix_token: "<|begin_of_text|>"
     end_of_text_token: "<|end_of_text|>"
 
+  - name: meta/llama-3.1-8b
+    tokenizer_spec:
+      class_name: "helm.tokenizers.huggingface_tokenizer.HuggingFaceTokenizer"
+      args:
+        pretrained_model_name_or_path: meta-llama/Meta-Llama-3.1-8B-Instruct
+    prefix_token: "<|begin_of_text|>"
+    end_of_text_token: "<|end_of_text|>"
+
   # 01-ai
   - name: 01-ai/Yi-6B
     tokenizer_spec:
@@ -247,6 +332,17 @@ tokenizer_configs:
     end_of_text_token: "</s>"
     prefix_token: "<s>"
 
+  # AI Singapore
+  - name: aisingapore/sea-lion-7b
+    tokenizer_spec:
+      class_name: "helm.tokenizers.huggingface_tokenizer.HuggingFaceTokenizer"
+      args:
+        trust_remote_code: true
+        use_fast: false
+    end_of_text_token: "<|endoftext|>"
+    prefix_token: ""
+
+
 
   # Allen Institute for AI
   # The allenai/olmo-7b requires Python 3.9 or newer.
@@ -259,6 +355,12 @@ tokenizer_configs:
     end_of_text_token: "<|endoftext|>"
     prefix_token: ""
 
+  - name: allenai/OLMo-1.7-7B-hf
+    tokenizer_spec:
+      class_name: "helm.tokenizers.huggingface_tokenizer.HuggingFaceTokenizer"
+    end_of_text_token: "<|endoftext|>"
+    prefix_token: ""
+
 
   # Microsoft
   - name: microsoft/phi-2  
@@ -267,6 +369,20 @@ tokenizer_configs:
     end_of_text_token: "<|endoftext|>"  
     prefix_token: "<|endoftext|>"
 
+  - name: microsoft/phi-3-small-8k-instruct
+    tokenizer_spec:  
+      class_name: "helm.tokenizers.huggingface_tokenizer.HuggingFaceTokenizer"
+      args:
+        trust_remote_code: true
+    end_of_text_token: "<|endoftext|>"
+    prefix_token: "<|endoftext|>"
+
+  - name: microsoft/phi-3-medium-4k-instruct
+    tokenizer_spec:  
+      class_name: "helm.tokenizers.huggingface_tokenizer.HuggingFaceTokenizer"
+    end_of_text_token: "<|endoftext|>"
+    prefix_token: "<s>"
+
   # Mistralai
   - name: mistralai/Mistral-7B-v0.1
     tokenizer_spec:
@@ -274,6 +390,36 @@ tokenizer_configs:
     end_of_text_token: "</s>"
     prefix_token: "<s>"
 
+  - name: mistralai/Mistral-7B-Instruct-v0.1
+    tokenizer_spec:
+      class_name: "helm.tokenizers.huggingface_tokenizer.HuggingFaceTokenizer"
+    end_of_text_token: "</s>"
+    prefix_token: "<s>"
+
+  - name: mistralai/Mistral-7B-Instruct-v0.2
+    tokenizer_spec:
+      class_name: "helm.tokenizers.huggingface_tokenizer.HuggingFaceTokenizer"
+    end_of_text_token: "</s>"
+    prefix_token: "<s>"
+
+  - name: mistralai/Mistral-7B-Instruct-v0.3
+    tokenizer_spec:
+      class_name: "helm.tokenizers.huggingface_tokenizer.HuggingFaceTokenizer"
+    end_of_text_token: "</s>"
+    prefix_token: "<s>"
+
+  - name: mistralai/Mistral-Nemo-Base-2407
+    tokenizer_spec:
+      class_name: "helm.tokenizers.huggingface_tokenizer.HuggingFaceTokenizer"
+    end_of_text_token: "</s>"
+    prefix_token: "<s>"
+
+  - name: mistralai/Mistral-Large-Instruct-2407
+    tokenizer_spec:
+      class_name: "helm.tokenizers.huggingface_tokenizer.HuggingFaceTokenizer"
+    end_of_text_token: "</s>"
+    prefix_token: "<s>"
+
   # Neurips
   - name: neurips/local
     tokenizer_spec:
@@ -281,19 +427,43 @@ tokenizer_configs:
     end_of_text_token: "<|endoftext|>"
     prefix_token: "<|endoftext|>"
 
-  # Openai
+  # NVIDIA
+  - name: nvidia/nemotron-4-340b-instruct
+    tokenizer_spec:
+      class_name: "helm.tokenizers.huggingface_tokenizer.HuggingFaceTokenizer"
+      args:
+        pretrained_model_name_or_path: Xenova/Nemotron-4-340B-Instruct-Tokenizer
+        revision: b7aa0de92cda9f9e722d58d6ca90f46ae17d4701
+    end_of_text_token: "<|endoftext|>"
+    prefix_token: "<|endoftext|>"
+
+  # OpenAI
   - name: openai/cl100k_base
     tokenizer_spec:
       class_name: "helm.tokenizers.tiktoken_tokenizer.TiktokenTokenizer"
     end_of_text_token: "<|endoftext|>"
     prefix_token: "<|endoftext|>"
 
+  - name: openai/o200k_base
+    tokenizer_spec:
+      class_name: "helm.tokenizers.tiktoken_tokenizer.TiktokenTokenizer"
+    end_of_text_token: "<|endoftext|>"
+    prefix_token: "<|endoftext|>"
+
   - name: openai/clip-vit-large-patch14
     tokenizer_spec:
       class_name: "helm.tokenizers.huggingface_tokenizer.HuggingFaceTokenizer"
     end_of_text_token: ""
     prefix_token: ""
 
+  # OpenThaiGPT
+  - name: openthaigpt/openthaigpt-1.0.0-7b-chat
+    tokenizer_spec:
+      class_name: "helm.tokenizers.huggingface_tokenizer.HuggingFaceTokenizer"
+    end_of_text_token: "</s>"
+    prefix_token: "<s>"
+
+  # Qwen
   - name: qwen/qwen-7b
     tokenizer_spec:
       class_name: "helm.tokenizers.huggingface_tokenizer.HuggingFaceTokenizer"
@@ -311,6 +481,14 @@ tokenizer_configs:
     end_of_text_token: "<|endoftext|>"
     prefix_token: ""
 
+  - name: qwen/qwen2-72b-instruct
+    tokenizer_spec:
+      class_name: "helm.tokenizers.huggingface_tokenizer.HuggingFaceTokenizer"
+      args:
+        pretrained_model_name_or_path: Qwen/Qwen2-72B-Instruct
+    end_of_text_token: <|im_end|>"
+    prefix_token: "<|im_start|>'"
+
   - name: qwen/qwen-vl
     tokenizer_spec:
       class_name: "helm.tokenizers.huggingface_tokenizer.HuggingFaceTokenizer"
@@ -331,10 +509,22 @@ tokenizer_configs:
     end_of_text_token: "<|endoftext|>"
     prefix_token: ""
 
+  # SambaLingo
+  - name: sambanova/sambalingo-thai-base
+    tokenizer_spec:
+      class_name: "helm.tokenizers.huggingface_tokenizer.HuggingFaceTokenizer"
+      args:
+        pretrained_model_name_or_path: sambanovasystems/SambaLingo-Thai-Base
+    end_of_text_token: "</s>"
+    prefix_token: "<s>"
+
   # Snowflake
   - name: snowflake/snowflake-arctic-instruct
     tokenizer_spec:
       class_name: "helm.tokenizers.huggingface_tokenizer.HuggingFaceTokenizer"
+      args:
+        pretrained_model_name_or_path: Snowflake/snowflake-arctic-instruct
+        trust_remote_code: true
     end_of_text_token: "<|im_end|>"
     prefix_token: "<|im_start|>"
 
@@ -352,6 +542,13 @@ tokenizer_configs:
     end_of_text_token: "</s>"
     prefix_token: ""
 
+  # Typhoon
+  - name: scb10x/typhoon-7b
+    tokenizer_spec:
+      class_name: "helm.tokenizers.huggingface_tokenizer.HuggingFaceTokenizer"
+    end_of_text_token: "</s>"
+    prefix_token: "<s>"
+
   # Writer
   - name: writer/gpt2
     tokenizer_spec:

helm/proxy/critique/model_critique_client.py

@@ -15,6 +15,7 @@ from helm.common.optional_dependencies import handle_module_not_found_error
 from helm.common.request import Request, RequestResult, GeneratedOutput
 from helm.clients.client import Client
 from helm.proxy.critique.critique_client import CritiqueClient
+from helm.common.media_object import MultimediaObject, MediaObject
 
 
 class CritiqueParseError(Exception):
@@ -24,6 +25,8 @@ class CritiqueParseError(Exception):
 class ModelCritiqueClient(CritiqueClient):
     """A CritiqueClient that queries a Model to answer CritiqueRequests."""
 
+    VISION_LANGUAGE_MODELS = ["openai/gpt-4-vision", "reka/reka", "huggingface/prometheus-vision"]
+
     def __init__(self, client: Client, model_name):
         self._client = client
         self._model_name = model_name
@@ -31,6 +34,11 @@ class ModelCritiqueClient(CritiqueClient):
             get_default_model_deployment_for_model(model_name, warn_arg_deprecated=False, ignore_deprecated=True)
             or self._model_name
         )
+        self.vision_language = False
+        for vision_language_model_name in self.VISION_LANGUAGE_MODELS:
+            if model_name.startswith(vision_language_model_name):
+                self.vision_language = True
+                break
 
     def _interpolate_fields(self, text: str, fields: Dict[str, str]) -> str:
         for key, value in fields.items():
@@ -58,10 +66,15 @@ class ModelCritiqueClient(CritiqueClient):
 
         requests: List[Request] = []
         for question in task.questions:
-            prompt: str = base_prompt + "\n\n" + self._question_to_prompt(question, fields)
+            prompt: str
+            if len(question.text) > 0:
+                prompt = base_prompt + "\n\n" + self._question_to_prompt(question, fields)
+            else:
+                # We may don't want to add extra newlines and prompts
+                # if the question text is empty (e.g., the Vibe-Eval evaluator).
+                prompt = base_prompt
             if question.question_type == "free_response":
-                # TODO: Make max_tokens configurable
-                max_tokens = 100
+                max_tokens = 100 if task.max_tokens is None else task.max_tokens
             elif question.question_type == "checkbox":
                 # We multiply by 2 because the model will generate a comma after each option.
                 max_tokens = len(question.options) * 2
@@ -78,12 +91,21 @@ class ModelCritiqueClient(CritiqueClient):
 
                 prompt = anthropic.HUMAN_PROMPT + prompt + anthropic.AI_PROMPT
 
+            multimodal_prompt: Optional[MultimediaObject] = None
+            if self.vision_language:
+                assert question.media_object is not None, "Expect media_object for vision-language models"
+                image_media: MediaObject = question.media_object
+                text_media: MediaObject = MediaObject(text=prompt, content_type="text/plain")
+                multimodal_prompt = MultimediaObject(media_objects=[image_media, text_media])
+                prompt = ""  # set to empty string to avoid conflicts with multimodal_prompt
+
             request = Request(
                 model=self._model_name,
                 model_deployment=self._model_deployment_name,
                 prompt=prompt,
                 max_tokens=max_tokens,
                 echo_prompt=False,
+                multimodal_prompt=multimodal_prompt,
             )
             requests.append(request)
         return requests
@@ -124,7 +146,13 @@ class ModelCritiqueClient(CritiqueClient):
                 raise CritiqueParseError(
                     f"Invalid answer: {completion}. Multiple choice questions should have one answer."
                 )
-            return answers[0]
+            letter_answer = answers[0]
+            choice_rank = string.ascii_uppercase.index(letter_answer)
+            if choice_rank >= len(question.options):
+                raise CritiqueParseError(
+                    f"Invalid answer: {completion}. The answer is out of range of the options: {question.options}"
+                )
+            return letter_answer
         except CritiqueParseError as e:
             # If there was an error parsing the answer, we assume the user did not answer the question.
             hlog(f"Error parsing answer: {e}. Skipping question (and so the respondent entirely)")

helm/proxy/example_queries.py

@@ -22,7 +22,6 @@ example_queries = [
             temperature: 0.5  # Medium amount of randomness
             stop_sequences: [.]  # Stop when you hit a period
             model: openai/gpt-3.5-turbo-0613
-            model_deployment: openai/gpt-3.5-turbo-0613
             """
         ),
         environments="",
@@ -35,24 +34,24 @@ example_queries = [
             stop_sequences: [\\n]  # Stop when you hit a newline
             num_completions: 5  # Generate many samples
             model: openai/gpt-3.5-turbo-0613
-            model_deployment: openai/gpt-3.5-turbo-0613
             """
         ),
         environments="",
     ),
-    Query(
-        prompt="The quick brown fox jumps over the lazy dog.",
-        settings=dedent(
-            """
-            echo_prompt: true  # Analyze the prompt
-            max_tokens: 0  # Don't generate any more
-            top_k_per_token: 5  # Show alternatives for each position
-            model: openai/davinci-002
-            model_deployment: openai/davinci-002
-            """
-        ),
-        environments=dedent(""),
-    ),
+    # Disabled because `max_tokens: 0` no longer works on the OpenAI API
+    # Query(
+    #     prompt="The quick brown fox jumps over the lazy dog.",
+    #     settings=dedent(
+    #         """
+    #         echo_prompt: true  # Analyze the prompt
+    #         max_tokens: 0  # Don't generate any more
+    #         top_k_per_token: 5  # Show alternatives for each position
+    #         model: openai/text-davinci-002
+    #         model_deployment: openai/text-davinci-002
+    #         """
+    #     ),
+    #     environments=dedent(""),
+    # ),
     Query(
         prompt="Odd numbers: 1 -> 3 -> 5",
         settings=dedent(
@@ -60,7 +59,6 @@ example_queries = [
             temperature: 0  # Deterministic
             max_tokens: 50
             model: openai/gpt-3.5-turbo-0613
-            model_deployment: openai/gpt-3.5-turbo-0613
             """
         ),
         environments="",
@@ -73,7 +71,6 @@ example_queries = [
             stop_sequences: [.]
             # Try out multiple models
             model: ${model}
-            model_deployment: ${model}
             """
         ),
         environments=dedent(
@@ -100,7 +97,6 @@ example_queries = [
             num_completions: 5
             # Try out multiple models
             model: ${model}
-            model_deployment: ${model}
             """
         ),
         environments=dedent(
@@ -136,7 +132,6 @@ example_queries = [
             top_k_per_token: 4
             # Try out multiple models
             model: ${model}
-            model_deployment: ${model}
             """
         ),
         environments=dedent(
@@ -150,7 +145,6 @@ example_queries = [
         settings=dedent(
             """
             model: openai/gpt-3.5-turbo-0613
-            model_deployment: openai/gpt-3.5-turbo-0613
             """
         ),
         environments="",
@@ -163,7 +157,6 @@ example_queries = [
             stop_sequences: [\\n]
             # Try out multiple models
             model: ${model}
-            model_deployment: ${model}
             """
         ),
         environments=dedent(

helm/proxy/services/server_service.py

@@ -25,7 +25,6 @@ from helm.common.hierarchical_logger import hlog
 from helm.proxy.accounts import Accounts, Account
 from helm.clients.auto_client import AutoClient
 from helm.clients.moderation_api_client import ModerationAPIClient
-from helm.clients.perspective_api_client import PerspectiveAPIClient
 from helm.clients.image_generation.nudity_check_client import NudityCheckClient
 from helm.clients.gcs_client import GCSClient
 from helm.clients.clip_score_client import CLIPScoreClient
@@ -75,7 +74,7 @@ class ServerService(Service):
         # Lazily instantiate the following clients
         self.moderation_api_client: Optional[ModerationAPIClient] = None
         self.toxicity_classifier_client: Optional[ToxicityClassifierClient] = None
-        self.perspective_api_client: Optional[PerspectiveAPIClient] = None
+        self.perspective_api_client: Optional[ToxicityClassifierClient] = None
         self.nudity_check_client: Optional[NudityCheckClient] = None
         self.clip_score_client: Optional[CLIPScoreClient] = None
         self.gcs_client: Optional[GCSClient] = None
@@ -119,7 +118,7 @@ class ServerService(Service):
                 return "codex"
             elif model_deployment.startswith("openai/dall-e-"):
                 return "dall_e"
-            elif model_deployment.startswith("openai/gpt-4-"):
+            elif model_deployment.startswith("openai/gpt-4"):
                 return "gpt4"
             else:
                 return "gpt3"

helm/proxy/token_counters/test_auto_token_counter.py

@@ -13,8 +13,8 @@ class TestAutoTokenCounter:
         )
         # The following prompt has 51 tokens according to the GPT-2 tokenizer
         request = Request(
-            model="openai/text-davinci-002",
-            model_deployment="openai/text-davinci-002",
+            model="openai/gpt2",
+            model_deployment="huggingface/gpt2",
             prompt="The Center for Research on Foundation Models (CRFM) is "
             "an interdisciplinary initiative born out of the Stanford "
             "Institute for Human-Centered Artificial Intelligence (HAI) "

helm/tokenizers/ai21_tokenizer.py

@@ -1,60 +1,52 @@
-from typing import Any, Dict, List
-import requests
-
-from dacite import from_dict
-
-from helm.common.cache import Cache, CacheConfig
-from helm.common.tokenization_request import (
-    TokenizationRequest,
-    TokenizationRequestResult,
-    TokenizationToken,
-    TextRange,
-    DecodeRequest,
-    DecodeRequestResult,
-)
-from helm.clients.ai21_utils import AI21RequestError, handle_failed_request
-from .tokenizer import Tokenizer
-
-
-class AI21Tokenizer(Tokenizer):
-    def __init__(self, api_key: str, cache_config: CacheConfig) -> None:
-        self.cache = Cache(cache_config)
-        self.api_key: str = api_key
-
-    def tokenize(self, request: TokenizationRequest) -> TokenizationRequestResult:
-        """
-        Tokenizes the text by using the AI21 endpoint: https://api.ai21.com/studio/v1/tokenize.
-        """
-        # TODO: Does not support encoding
-        raw_request: Dict[str, str] = {"text": request.text}
-
-        def do_it() -> Dict[str, Any]:
-            response = requests.post(
-                "https://api.ai21.com/studio/v1/tokenize",
-                headers={"Authorization": f"Bearer {self.api_key}"},
-                json=raw_request,
-            ).json()
-
-            # If 'tokens' is not present in the response, assume request failed.
-            if "tokens" not in response:
-                handle_failed_request(api_type="tokenizer", response=response)
-
-            return response
-
-        try:
-            response, cached = self.cache.get(raw_request, do_it)
-        except AI21RequestError:
-            return TokenizationRequestResult(success=False, cached=False, text="", tokens=[])
-
-        # Each token is represented like this in the response:
-        # {'token': '▁Hello', 'textRange': {'start': 0, 'end': 5}}
-        tokens: List[TokenizationToken] = []
-        for token_dict in response["tokens"]:
-            tokens.append(
-                TokenizationToken(value=token_dict["token"], text_range=from_dict(TextRange, token_dict["textRange"]))
+import threading
+from typing import Any, Dict
+
+from helm.common.cache import CacheConfig
+from helm.common.optional_dependencies import handle_module_not_found_error
+from helm.tokenizers.caching_tokenizer import CachingTokenizer
+
+try:
+    from ai21_tokenizer import Tokenizer as SDKTokenizer
+    from ai21_tokenizer.base_tokenizer import BaseTokenizer
+except ModuleNotFoundError as e:
+    handle_module_not_found_error(e, ["ai21"])
+
+
+class AI21LocalTokenizer(CachingTokenizer):
+    """AI21 tokenizer using the AI21 Python library."""
+
+    def __init__(self, cache_config: CacheConfig) -> None:
+        super().__init__(cache_config)
+        self._tokenizers_lock = threading.Lock()
+        self.tokenizers: Dict[str, BaseTokenizer] = {}
+
+    def _get_tokenizer(self, tokenizer_name: str) -> BaseTokenizer:
+        with self._tokenizers_lock:
+            if tokenizer_name not in self.tokenizers:
+                self.tokenizers[tokenizer_name] = SDKTokenizer.get_tokenizer(tokenizer_name)
+            return self.tokenizers[tokenizer_name]
+
+    def _tokenize_do_it(self, request: Dict[str, Any]) -> Dict[str, Any]:
+        tokenizer_name = request["tokenizer"].split("/")[1]
+        tokenizer = self._get_tokenizer(tokenizer_name)
+        if request["truncation"]:
+            token_ids = tokenizer.encode(
+                text=request["text"],
+                truncation=request["truncation"],
+                max_length=request["max_length"],
+                add_special_tokens=False,
             )
-        text: str = response["text"]
-        return TokenizationRequestResult(success=True, cached=cached, tokens=tokens, text=text)
-
-    def decode(self, request: DecodeRequest) -> DecodeRequestResult:
-        raise NotImplementedError("Not supported")
+        else:
+            token_ids = tokenizer.encode(
+                text=request["text"],
+                add_special_tokens=False,
+            )
+        if request["encode"]:
+            return {"tokens": token_ids}
+        else:
+            return {"tokens": tokenizer.convert_ids_to_tokens(token_ids)}
+
+    def _decode_do_it(self, request: Dict[str, Any]) -> Dict[str, Any]:
+        tokenizer_name = request["tokenizer"].split("/")[1]
+        tokenizer = self._get_tokenizer(tokenizer_name)
+        return {"text": tokenizer.decode(request["tokens"])}

helm/tokenizers/auto_tokenizer.py

@@ -41,7 +41,7 @@ class AutoTokenizer(Tokenizer):
         if tokenizer_config:
             tokenizer_spec = inject_object_spec_args(
                 tokenizer_config.tokenizer_spec,
-                constant_bindings={"cache_config": cache_config},
+                constant_bindings={"cache_config": cache_config, "tokenizer_name": tokenizer_name},
                 provider_bindings={
                     "api_key": lambda: provide_api_key(self.credentials, organization),
                     "project_id": lambda: self.credentials.get(organization + "ProjectId", None),  # VertexAI