crfm-helm 0.2.0__py3-none-any.whl → 0.2.2__py3-none-any.whl

helm/common/codec.py

@@ -0,0 +1,123 @@
+"""Functions for converting to and from dataclasses."""
+
+import dataclasses
+import json
+import typing
+from typing import Any, Callable, Dict, List, Union, Type, TypeVar
+
+from helm.benchmark.augmentations.dialect_perturbation import DialectPerturbation
+from helm.benchmark.augmentations.extra_space_perturbation import ExtraSpacePerturbation
+from helm.benchmark.augmentations.filler_words_perturbation import FillerWordsPerturbation
+from helm.benchmark.augmentations.gender_perturbation import GenderPerturbation
+from helm.benchmark.augmentations.misspelling_perturbation import MisspellingPerturbation
+from helm.benchmark.augmentations.person_name_perturbation import PersonNamePerturbation
+from helm.benchmark.augmentations.space_perturbation import SpacePerturbation
+from helm.benchmark.augmentations.synonym_perturbation import SynonymPerturbation
+from helm.benchmark.augmentations.typos_perturbation import TyposPerturbation
+from helm.benchmark.augmentations.perturbation_description import PerturbationDescription
+
+import cattrs
+from cattrs.gen import make_dict_structure_fn, make_dict_unstructure_fn
+
+
+T = TypeVar("T")
+StructureFn = Callable[[Dict[str, Any], Type[T]], T]  # dict -> dataclass
+UnstructureFn = Callable[[T], Dict[str, Any]]  # dataclass -> dict
+
+
+# TODO(#1251): Add proper class registration
+PERTURBATION_NAME_TO_DESCRIPTION = {
+    DialectPerturbation.name: DialectPerturbation.Description,
+    ExtraSpacePerturbation.name: ExtraSpacePerturbation.Description,
+    FillerWordsPerturbation.name: FillerWordsPerturbation.Description,
+    GenderPerturbation.name: GenderPerturbation.Description,
+    MisspellingPerturbation.name: MisspellingPerturbation.Description,
+    PersonNamePerturbation.name: PersonNamePerturbation.Description,
+    SpacePerturbation.name: SpacePerturbation.Description,
+    SynonymPerturbation.name: SynonymPerturbation.Description,
+    TyposPerturbation.name: TyposPerturbation.Description,
+}
+
+
+def _build_converter() -> cattrs.Converter:
+    converter = cattrs.Converter()
+
+    # Handle omission of Nones in JSON.
+    # To improve readability and reduce storage space, if a field value is None and the field
+    # has no default value or a None default value, the field is omitted in the serialized JSON.
+    def get_dataclass_optional_fields_without_default(cls: Type[T]) -> List[str]:
+        if not dataclasses.is_dataclass(cls):
+            return []
+        return [
+            field.name
+            for field in dataclasses.fields(cls)
+            if typing.get_origin(field.type) == Union and type(None) in typing.get_args(field.type)
+            # For optional fields with a non-None default value, do not replace a missing value
+            # with None.
+            and (field.default == dataclasses.MISSING or field.default is None)
+            and field.default_factory == dataclasses.MISSING
+        ]
+
+    def make_omit_nones_dict_structure_fn(cls: Type[T]) -> StructureFn[T]:
+        field_names = get_dataclass_optional_fields_without_default(cls)
+        _base_structure = make_dict_structure_fn(cls, converter)
+
+        def structure(raw_dict: Dict[str, Any], inner_cls: Type[T]) -> T:
+            for field_name in field_names:
+                if field_name not in raw_dict:
+                    raw_dict[field_name] = None
+            return _base_structure(raw_dict, inner_cls)
+
+        return structure
+
+    def make_omit_nones_dict_unstructure_fn(cls: Type[T]) -> UnstructureFn[T]:
+        field_names = get_dataclass_optional_fields_without_default(cls)
+        _base_unstructure = make_dict_unstructure_fn(cls, converter)
+
+        def structure(data: T) -> Dict[str, Any]:
+            raw_dict = _base_unstructure(data)
+            for field_name in field_names:
+                if raw_dict[field_name] is None:
+                    del raw_dict[field_name]
+            return raw_dict
+
+        return structure
+
+    converter.register_structure_hook_factory(
+        lambda cls: bool(get_dataclass_optional_fields_without_default(cls)), make_omit_nones_dict_structure_fn
+    )
+    converter.register_unstructure_hook_factory(
+        lambda cls: bool(get_dataclass_optional_fields_without_default(cls)), make_omit_nones_dict_unstructure_fn
+    )
+
+    # Handle the use of the name field in PerturbationDescription to determine the subclass.
+    base_perturbation_description_structure_fn: StructureFn = make_omit_nones_dict_structure_fn(PerturbationDescription)
+    perturbation_name_to_base_structure_fn: Dict[str, StructureFn] = {
+        name: make_omit_nones_dict_structure_fn(cls) for name, cls in PERTURBATION_NAME_TO_DESCRIPTION.items()
+    }
+
+    def structure_perturbation_description(
+        raw_dict: Dict[Any, Any], cls: Type[PerturbationDescription]
+    ) -> PerturbationDescription:
+        """Convert a raw dictionary to a PerturbationDescription.
+        This uses the name field to look up the correct PerturbationDescription subclass to output.
+        """
+        structure = perturbation_name_to_base_structure_fn.get(
+            raw_dict["name"], base_perturbation_description_structure_fn
+        )
+        return structure(raw_dict, cls)
+
+    converter.register_structure_hook(PerturbationDescription, structure_perturbation_description)
+
+    return converter
+
+
+_converter = _build_converter()
+
+
+def from_json(data: Union[bytes, str], cls: Type[T]) -> T:
+    return _converter.structure(json.loads(data), cls)
+
+
+def to_json(data: Any) -> str:
+    return json.dumps(_converter.unstructure(data), indent=2)

helm/common/general.py

@@ -49,7 +49,13 @@ def shell(args: List[str]):
 
 
 @htrack(None)
-def ensure_file_downloaded(source_url: str, target_path: str, unpack: bool = False, unpack_type: Optional[str] = None):
+def ensure_file_downloaded(
+    source_url: str,
+    target_path: str,
+    unpack: bool = False,
+    downloader_executable: str = "wget",
+    unpack_type: Optional[str] = None,
+):
     """Download `source_url` to `target_path` if it doesn't exist."""
     if os.path.exists(target_path):
         # Assume it's all good
@@ -59,7 +65,8 @@ def ensure_file_downloaded(source_url: str, target_path: str, unpack: bool = Fal
     # Download
     # gdown is used to download large files/zip folders from Google Drive.
     # It bypasses security warnings which wget cannot handle.
-    downloader_executable: str = "gdown" if source_url.startswith("https://drive.google.com") else "wget"
+    if source_url.startswith("https://drive.google.com"):
+        downloader_executable = "gdown"
     tmp_path: str = f"{target_path}.tmp"
     shell([downloader_executable, source_url, "-O", tmp_path])
 
@@ -195,13 +202,13 @@ def parallel_map(
     with htrack_block(f"Parallelizing computation on {len(items)} items over {parallelism} {units}"):
         results: List
         if parallelism == 1:
-            results = list(tqdm(map(process, items), total=len(items)))
+            results = list(tqdm(map(process, items), total=len(items), disable=None))
         elif multiprocessing:
             with ProcessPoolExecutor(max_workers=parallelism) as executor:
-                results = list(tqdm(executor.map(process, items), total=len(items)))
+                results = list(tqdm(executor.map(process, items), total=len(items), disable=None))
         else:
             with ThreadPoolExecutor(max_workers=parallelism) as executor:
-                results = list(tqdm(executor.map(process, items), total=len(items)))
+                results = list(tqdm(executor.map(process, items), total=len(items), disable=None))
     return results
 
 

helm/common/test_codec.py

@@ -0,0 +1,144 @@
+import unittest
+import json
+
+from dataclasses import dataclass
+
+from typing import Dict, List, Optional
+
+from helm.benchmark.augmentations.dialect_perturbation import DialectPerturbation
+from helm.benchmark.augmentations.extra_space_perturbation import ExtraSpacePerturbation
+from helm.benchmark.augmentations.filler_words_perturbation import FillerWordsPerturbation
+from helm.benchmark.augmentations.gender_perturbation import GenderPerturbation
+from helm.benchmark.augmentations.misspelling_perturbation import MisspellingPerturbation
+from helm.benchmark.augmentations.person_name_perturbation import PersonNamePerturbation
+from helm.benchmark.augmentations.space_perturbation import SpacePerturbation
+from helm.benchmark.augmentations.synonym_perturbation import SynonymPerturbation
+from helm.benchmark.augmentations.typos_perturbation import TyposPerturbation
+from helm.benchmark.augmentations.perturbation_description import PerturbationDescription
+from helm.common.codec import from_json, to_json
+
+
+@dataclass(frozen=True)
+class DataClassChildForTest:
+    required_int: int
+
+
+@dataclass(frozen=True)
+class DataClassWithOptionals:
+    optional_str: Optional[str]
+    optional_int: Optional[int]
+    optional_bool: Optional[bool]
+    optional_list: Optional[List[int]]
+    optional_dict: Optional[Dict[str, int]]
+    optional_child: Optional[DataClassChildForTest]
+
+
+@dataclass(frozen=True)
+class DataClassWithDefaults:
+    required_int_with_default: int = -1
+    optional_int_with_int_default: Optional[int] = -2
+    optional_int_with_none_default: Optional[int] = None
+
+
+class TestJsonCodec(unittest.TestCase):
+    def test_round_trip_optional(self):
+        data = DataClassWithOptionals(
+            optional_str="hello",
+            optional_int=42,
+            optional_bool=True,
+            optional_list=[2, 3, 5],
+            optional_dict={"x": 7},
+            optional_child=DataClassChildForTest(137),
+        )
+        self.assertEqual(data, from_json(to_json(data), DataClassWithOptionals))
+
+    def test_round_trip_optional_nones(self):
+        data = DataClassWithOptionals(
+            optional_str=None,
+            optional_int=None,
+            optional_bool=None,
+            optional_list=None,
+            optional_dict=None,
+            optional_child=None,
+        )
+        data_json = to_json(data)
+        self.assertEqual("{}", data_json)
+        self.assertEqual(data, from_json(data_json, DataClassWithOptionals))
+
+    def test_round_trip_default(self):
+        data = DataClassWithDefaults()
+        data_json = to_json(data)
+        self.assertCountEqual(
+            {"required_int_with_default": -1, "optional_int_with_int_default": -2}.items(),
+            json.loads(data_json).items(),
+        )
+        self.assertEqual(data, from_json(data_json, DataClassWithDefaults))
+
+    def test_round_trip_default_ints(self):
+        data = DataClassWithDefaults(
+            required_int_with_default=1,
+            optional_int_with_int_default=2,
+            optional_int_with_none_default=3,
+        )
+        data_json = to_json(data)
+        self.assertEqual(data, from_json(data_json, DataClassWithDefaults))
+
+    def test_round_trip_default_nones(self):
+        data = DataClassWithDefaults(
+            optional_int_with_int_default=None,
+            optional_int_with_none_default=None,
+        )
+        data_json = to_json(data)
+        self.assertCountEqual(
+            {
+                "required_int_with_default": -1,
+                # `optional_int_with_int_default` should deserialize back to None,
+                # rather than the default int value. Therefore it must be
+                # serialized to null in JSON instead of removed.
+                "optional_int_with_int_default": None,
+            }.items(),
+            json.loads(data_json).items(),
+        )
+        self.assertEqual(data, from_json(data_json, DataClassWithDefaults))
+
+    def test_round_trip_perturbation_descriptions(self):
+        descriptions = [
+            PerturbationDescription(
+                name="unknown",
+            ),
+            DialectPerturbation.Description(
+                name=DialectPerturbation.name,
+                fairness=True,
+                prob=0.5,
+                source_class="source_class",
+                target_class="target_class",
+                mapping_file_path="mapping_file_path",
+            ),
+            ExtraSpacePerturbation.Description(name=ExtraSpacePerturbation.name, robustness=True, num_spaces=2),
+            FillerWordsPerturbation.Description(name=FillerWordsPerturbation.name, robustness=True, insert_prob=0.5),
+            GenderPerturbation.Description(
+                name=GenderPerturbation.name,
+                mode="mode",
+                fairness=True,
+                prob=0.5,
+                source_class="source_class",
+                target_class="target_class",
+                bidirectional=True,
+            ),
+            MisspellingPerturbation.Description(name=MisspellingPerturbation.name, robustness=True, prob=0.5),
+            PersonNamePerturbation.Description(
+                name=PersonNamePerturbation.name,
+                fairness=True,
+                prob=0.5,
+                source_class="source_str",
+                target_class="target_str",
+                name_file_path="name_file_path",
+                person_name_type="person_name_type",
+                preserve_gender=True,
+            ),
+            SpacePerturbation.Description(name=SpacePerturbation.name, robustness=True, max_spaces=2),
+            SynonymPerturbation.Description(name=SynonymPerturbation.name, robustness=True, prob=0.5),
+            TyposPerturbation.Description(name=TyposPerturbation.name, robustness=True, prob=0.5),
+        ]
+        for description in descriptions:
+            self.assertEqual(description, from_json(to_json(description), PerturbationDescription))

helm/proxy/clients/aleph_alpha_client.py

@@ -2,7 +2,11 @@ import json
 import requests
 from typing import Any, Dict, List
 
+from aleph_alpha_client import Client as AlephAlphaPythonClient
+from tokenizers import Tokenizer, Encoding
+
 from helm.common.cache import Cache, CacheConfig
+from helm.common.hierarchical_logger import hlog
 from helm.common.request import Request, RequestResult, Sequence, Token
 from helm.common.tokenization_request import (
     DecodeRequest,
@@ -19,9 +23,27 @@ class AlephAlphaClient(Client):
     TOKENIZE_ENDPOINT: str = "tokenize"
     DETOKENIZE_ENDPOINT: str = "detokenize"
 
+    VALID_TOKENIZERS: List[str] = [
+        "luminous-base",
+        "luminous-extended",
+        "luminous-supreme",
+    ]
+
     def __init__(self, api_key: str, cache_config: CacheConfig):
         self.api_key: str = api_key
         self.cache = Cache(cache_config)
+        self._aleph_alpha_client = AlephAlphaPythonClient(token=api_key)
+        self._tokenizer_name_to_tokenizer: Dict[str, Tokenizer] = {}
+
+    def _get_tokenizer(self, tokenizer_name: str) -> Tokenizer:
+        if tokenizer_name not in self.VALID_TOKENIZERS:
+            raise ValueError(f"Invalid tokenizer: {tokenizer_name}")
+
+        # Check if the tokenizer is cached
+        if tokenizer_name not in self._tokenizer_name_to_tokenizer:
+            self._tokenizer_name_to_tokenizer[tokenizer_name] = self._aleph_alpha_client.tokenizer(tokenizer_name)
+            hlog(f"Initialized tokenizer: {tokenizer_name}")
+        return self._tokenizer_name_to_tokenizer[tokenizer_name]
 
     def _send_request(self, endpoint: str, raw_request: Dict[str, Any]) -> Dict[str, Any]:
         response = requests.request(
@@ -33,6 +55,8 @@ class AlephAlphaClient(Client):
                 "Authorization": f"Bearer {self.api_key}",
             },
             data=json.dumps(raw_request),
+            # Setting the nice flag prevents intensive benchmarking runs from saturating Aleph Alpha's API queues
+            params=json.dumps({"nice": True}),
         )
         result = json.loads(response.text)
         assert "error" not in result, f"Request failed with error: {result['error']}"
@@ -40,7 +64,6 @@ class AlephAlphaClient(Client):
 
     def make_request(self, request: Request) -> RequestResult:
         """Make a request following https://docs.aleph-alpha.com/api/complete."""
-        # TODO: echo is not supported. Follow up on this.
         raw_request = {
             "model": request.model_engine,
             "prompt": request.prompt,
@@ -53,6 +76,7 @@ class AlephAlphaClient(Client):
             "n": request.num_completions,
             "stop_sequences": request.stop_sequences,
             "log_probs": request.top_k_per_token,
+            "echo": request.echo_prompt,
             "tokens": True,  # Setting to True returns individual tokens of the completion
         }
 
@@ -102,24 +126,21 @@ class AlephAlphaClient(Client):
         )
 
     def tokenize(self, request: TokenizationRequest) -> TokenizationRequestResult:
-        """Make a request following https://docs.aleph-alpha.com/api/tokenize."""
-        raw_request = {
-            "model": request.tokenizer_name,
-            "prompt": request.text,
-            "tokens": True,
-            "token_ids": True,
-        }
-
+        """
+        Encode the text using Aleph Alpha's tokenizer library:
+        https://aleph-alpha-client.readthedocs.io/en/latest/aleph_alpha_client.html#aleph_alpha_client.Client.tokenizer
+        """
         try:
 
             def do_it():
-                result = self._send_request(AlephAlphaClient.TOKENIZE_ENDPOINT, raw_request)
-                assert "tokens" in result and "token_ids" in result, f"Invalid response: {result}"
-                return result
-
-            response, cached = self.cache.get(raw_request, wrap_request_time(do_it))
-        except (requests.exceptions.RequestException, AssertionError) as e:
-            error: str = f"AlephAlphaClient error: {e}"
+                tokenizer: Tokenizer = self._get_tokenizer(request.tokenizer_name)
+                result: Encoding = tokenizer.encode(request.text, add_special_tokens=False)
+                return {"token_ids": result.ids, "tokens": result.tokens}
+
+            cache_key = {"model": request.tokenizer_name, "prompt": request.text, "tokens": True, "token_ids": True}
+            response, cached = self.cache.get(cache_key, wrap_request_time(do_it))
+        except RuntimeError as e:
+            error: str = f"AlephAlphaClient tokenize error: {e}"
             return TokenizationRequestResult(error=error, success=False, cached=False, text="", tokens=[])
 
         tokens = response["token_ids" if request.encode else "tokens"]
@@ -135,22 +156,20 @@ class AlephAlphaClient(Client):
         )
 
     def decode(self, request: DecodeRequest) -> DecodeRequestResult:
-        """Make a request following https://docs.aleph-alpha.com/api/detokenize."""
-        raw_request = {
-            "model": request.tokenizer_name,
-            "token_ids": request.tokens,
-        }
-
+        """
+        Decode the tokens using Aleph Alpha's tokenizer library:
+        https://aleph-alpha-client.readthedocs.io/en/latest/aleph_alpha_client.html#aleph_alpha_client.Client.tokenizer
+        """
         try:
 
             def do_it():
-                result = self._send_request(AlephAlphaClient.DETOKENIZE_ENDPOINT, raw_request)
-                assert "result" in result, f"Invalid response: {result}"
-                return result
+                tokenizer: Tokenizer = self._get_tokenizer(request.tokenizer_name)
+                return {"result": tokenizer.decode(request.tokens)}
 
-            response, cached = self.cache.get(raw_request, wrap_request_time(do_it))
-        except (requests.exceptions.RequestException, AssertionError) as e:
-            error: str = f"AlephAlphaClient error: {e}"
+            cache_key = {"model": request.tokenizer_name, "token_ids": request.tokens}
+            response, cached = self.cache.get(cache_key, wrap_request_time(do_it))
+        except RuntimeError as e:
+            error: str = f"AlephAlphaClient decode error: {e}"
             return DecodeRequestResult(error=error, success=False, cached=False, text="")
 
         return DecodeRequestResult(

helm/proxy/clients/auto_client.py

@@ -21,6 +21,7 @@ from .anthropic_client import AnthropicClient
 from .chat_gpt_client import ChatGPTClient
 from .cohere_client import CohereClient
 from .together_client import TogetherClient
+from .google_client import GoogleClient
 from .goose_ai_client import GooseAIClient
 from .huggingface_client import HuggingFaceClient
 from .ice_tokenizer_client import ICETokenizerClient
@@ -29,6 +30,7 @@ from .microsoft_client import MicrosoftClient
 from .perspective_api_client import PerspectiveAPIClient
 from .yalm_tokenizer_client import YaLMTokenizerClient
 from .simple_client import SimpleClient
+from helm.proxy.clients.huggingface_model_registry import get_huggingface_model_config
 
 
 class AutoClient(Client):
@@ -53,15 +55,17 @@ class AutoClient(Client):
         # TODO: Allow setting CacheConfig.follower_cache_path from a command line flag.
         return SqliteCacheConfig(client_cache_path)
 
-    def get_client(self, request: Request) -> Client:
-        """Return a client based on `organization`, creating it if necessary."""
-        organization: str = request.model_organization
-        client: Optional[Client] = self.clients.get(organization)
+    def _get_client(self, model: str) -> Client:
+        """Return a client based on the model, creating it if necessary."""
+        client: Optional[Client] = self.clients.get(model)
 
         if client is None:
+            organization: str = model.split("/")[0]
             cache_config: CacheConfig = self._build_cache_config(organization)
 
-            if organization == "openai":
+            if get_huggingface_model_config(model):
+                client = HuggingFaceClient(cache_config=cache_config)
+            elif organization == "openai":
                 # TODO: add ChatGPT to the OpenAIClient when it's supported.
                 #       We're using a separate client for now since we're using an unofficial Python library.
                 # See https://github.com/acheong08/ChatGPT/wiki/Setup on how to get a valid session token.
@@ -71,13 +75,14 @@ class AutoClient(Client):
                     # TODO: use `cache_config` above. Since this feature is still experimental,
                     #       save queries and responses in a separate collection.
                     cache_config=self._build_cache_config("ChatGPT"),
-                    tokenizer_client=self.get_tokenizer_client("huggingface"),
+                    tokenizer_client=self._get_tokenizer_client("huggingface"),
                 )
 
                 org_id = self.credentials.get("openaiOrgId", None)
                 client = OpenAIClient(
                     api_key=self.credentials["openaiApiKey"],
                     cache_config=cache_config,
+                    tokenizer_client=self._get_tokenizer_client("huggingface"),
                     chat_gpt_client=chat_gpt_client,
                     org_id=org_id,
                 )
@@ -105,18 +110,20 @@ class AutoClient(Client):
                     cache_config=cache_config,
                     org_id=org_id,
                 )
+            elif organization == "google":
+                client = GoogleClient(cache_config=cache_config)
             elif organization == "together":
                 client = TogetherClient(api_key=self.credentials.get("togetherApiKey", None), cache_config=cache_config)
             elif organization == "simple":
                 client = SimpleClient(cache_config=cache_config)
             else:
-                raise ValueError(f"Unknown organization: {organization}")
-            self.clients[organization] = client
+                raise ValueError(f"Could not find client for model: {model}")
+            self.clients[model] = client
         return client
 
     def make_request(self, request: Request) -> RequestResult:
         """
-        Dispatch based on the organization in the name of the model (e.g., openai/davinci).
+        Dispatch based on the the name of the model (e.g., openai/davinci).
         Retries if request fails.
         """
 
@@ -125,30 +132,33 @@ class AutoClient(Client):
         def make_request_with_retry(client: Client, request: Request) -> RequestResult:
             return client.make_request(request)
 
-        organization: str = request.model_organization
-        client: Client = self.get_client(request)
+        client: Client = self._get_client(request.model)
 
         try:
             return make_request_with_retry(client=client, request=request)
         except RetryError as e:
             last_attempt: Attempt = e.last_attempt
             retry_error: str = (
-                f"Failed to make request to {organization} after retrying {last_attempt.attempt_number} times"
+                f"Failed to make request to {request.model} after retrying {last_attempt.attempt_number} times"
             )
             hlog(retry_error)
 
             # Notify our user that we failed to make the request even after retrying.
             return replace(last_attempt.value, error=f"{retry_error}. Error: {last_attempt.value.error}")
 
-    def get_tokenizer_client(self, organization: str) -> Client:
-        """Return a client based on `organization`, creating it if necessary."""
-        client: Optional[Client] = self.tokenizer_clients.get(organization)
+    def _get_tokenizer_client(self, tokenizer: str) -> Client:
+        """Return a client based on the tokenizer, creating it if necessary."""
+        organization: str = tokenizer.split("/")[0]
+        client: Optional[Client] = self.tokenizer_clients.get(tokenizer)
 
         if client is None:
             cache_config: CacheConfig = self._build_cache_config(organization)
-            if organization in [
+            if get_huggingface_model_config(tokenizer):
+                client = HuggingFaceClient(cache_config=cache_config)
+            elif organization in [
                 "anthropic",
                 "bigscience",
+                "bigcode",
                 "EleutherAI",
                 "facebook",
                 "google",
@@ -171,19 +181,18 @@ class AutoClient(Client):
             elif organization == "simple":
                 client = SimpleClient(cache_config=cache_config)
             else:
-                raise ValueError(f"Unknown organization: {organization}")
-            self.tokenizer_clients[organization] = client
+                raise ValueError(f"Could not find tokenizer client for model: {tokenizer}")
+            self.tokenizer_clients[tokenizer] = client
         return client
 
     def tokenize(self, request: TokenizationRequest) -> TokenizationRequestResult:
-        """Tokenizes based on the organization in the name of the tokenizer (e.g., huggingface/gpt2)."""
+        """Tokenizes based on the name of the tokenizer (e.g., huggingface/gpt2)."""
 
         @retry_request
         def tokenize_with_retry(client: Client, request: TokenizationRequest) -> TokenizationRequestResult:
             return client.tokenize(request)
 
-        organization: str = request.tokenizer_organization
-        client: Client = self.get_tokenizer_client(organization)
+        client: Client = self._get_tokenizer_client(request.tokenizer)
 
         try:
             return tokenize_with_retry(client=client, request=request)
@@ -194,14 +203,13 @@ class AutoClient(Client):
             return replace(last_attempt.value, error=f"{retry_error}. Error: {last_attempt.value.error}")
 
     def decode(self, request: DecodeRequest) -> DecodeRequestResult:
-        """Decodes based on the organization in the name of the tokenizer (e.g., huggingface/gpt2)."""
+        """Decodes based on the the name of the tokenizer (e.g., huggingface/gpt2)."""
 
         @retry_request
         def decode_with_retry(client: Client, request: DecodeRequest) -> DecodeRequestResult:
             return client.decode(request)
 
-        organization: str = request.tokenizer_organization
-        client: Client = self.get_tokenizer_client(organization)
+        client: Client = self._get_tokenizer_client(request.tokenizer)
 
         try:
             return decode_with_retry(client=client, request=request)