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

helm/proxy/clients/google_client.py

@@ -0,0 +1,88 @@
+from typing import List, Dict
+
+from helm.common.cache import Cache, CacheConfig
+from helm.common.request import Request, RequestResult, Sequence, Token
+from helm.common.tokenization_request import (
+    TokenizationRequest,
+    TokenizationRequestResult,
+    DecodeRequest,
+    DecodeRequestResult,
+)
+from .client import Client, truncate_sequence
+
+
+class GoogleClient(Client):
+    """
+    Client for the Google models. There isn't an API for their language models.
+    We receive and process completions offline.
+    """
+
+    @staticmethod
+    def convert_to_raw_request(request: Request) -> Dict:
+        return {
+            "best_of": request.top_k_per_token,
+            "echo": request.echo_prompt,
+            "logprobs": request.top_k_per_token,
+            "max_tokens": request.max_tokens,
+            "model": request.model_engine,
+            "n": request.num_completions,
+            "prompt": request.prompt,
+            "request_type": "language-model-inference",
+            "stop": request.stop_sequences or None,
+            "temperature": request.temperature,
+            "top_p": request.top_p,
+        }
+
+    def __init__(self, cache_config: CacheConfig):
+        self.cache = Cache(cache_config)
+
+    def make_request(self, request: Request) -> RequestResult:
+        raw_request = GoogleClient.convert_to_raw_request(request)
+        cache_key: Dict = Client.make_cache_key(raw_request, request)
+
+        try:
+
+            def fail():
+                raise RuntimeError(
+                    f"The result has not been uploaded to the cache for the following request: {cache_key}"
+                )
+
+            # If results are not cached for a given query, fail fast
+            response, cached = self.cache.get(cache_key, fail)
+        except RuntimeError as e:
+            error: str = f"GoogleClient error: {e}"
+            return RequestResult(success=False, cached=False, error=error, completions=[], embedding=[])
+
+        # Expect the result to be structured the same way as a response from OpenAI API.
+        completions: List[Sequence] = []
+        for raw_completion in response["choices"]:
+            sequence_logprob = 0
+            tokens: List[Token] = []
+
+            raw_data = raw_completion["logprobs"]
+            for text, logprob in zip(raw_data["tokens"], raw_data["token_logprobs"]):
+                tokens.append(Token(text=text, logprob=logprob or 0, top_logprobs=dict()))
+                sequence_logprob += logprob or 0
+
+            completion = Sequence(
+                text=raw_completion["text"],
+                logprob=sequence_logprob,
+                tokens=tokens,
+                finish_reason={"reason": raw_completion["finish_reason"]},
+            )
+            completion = truncate_sequence(completion, request)
+            completions.append(completion)
+
+        return RequestResult(
+            success=True,
+            cached=cached,
+            request_time=response["request_time"],
+            completions=completions,
+            embedding=[],
+        )
+
+    def tokenize(self, request: TokenizationRequest) -> TokenizationRequestResult:
+        raise NotImplementedError("Use the HuggingFaceClient to tokenize.")
+
+    def decode(self, request: DecodeRequest) -> DecodeRequestResult:
+        raise NotImplementedError("Use the HuggingFaceClient to decode.")

helm/proxy/clients/huggingface_client.py

@@ -1,3 +1,4 @@
+from copy import deepcopy
 import torch
 from dataclasses import asdict
 from transformers import AutoModelForCausalLM, AutoTokenizer
@@ -15,24 +16,29 @@ from helm.common.tokenization_request import (
 )
 from .client import Client, wrap_request_time, truncate_sequence
 from .huggingface_tokenizer import HuggingFaceTokenizers
+from helm.proxy.clients.huggingface_model_registry import HuggingFaceModelConfig, get_huggingface_model_config
 
 
 class HuggingFaceServer:
-    def __init__(self, model_name: str):
+    def __init__(self, model_config: HuggingFaceModelConfig):
         if torch.cuda.is_available():
             hlog("CUDA is available, initializing with a GPU...")
             self.device: str = "cuda:0"
         else:
             self.device = "cpu"
-
-        with htrack_block("Loading model"):
-            self.model = AutoModelForCausalLM.from_pretrained(model_name).to(self.device)
-        with htrack_block("Loading tokenizer"):
-            self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        model_kwargs = {}
+        if model_config.revision:
+            model_kwargs["revision"] = model_config.revision
+        with htrack_block(f"Loading Hugging Face model for config {model_config}"):
+            self.model = AutoModelForCausalLM.from_pretrained(
+                model_config.model_id, trust_remote_code=True, **model_kwargs
+            ).to(self.device)
+        with htrack_block(f"Loading Hugging Face tokenizer model for config {model_config}"):
+            self.tokenizer = AutoTokenizer.from_pretrained(model_config.model_id, **model_kwargs)
 
     def serve_request(self, raw_request: Dict[str, Any]):
         encoded_input = self.tokenizer(raw_request["prompt"], return_tensors="pt").to(self.device)
-
+        raw_request = deepcopy(raw_request)
         raw_request["do_sample"] = True
         raw_request["return_dict_in_generate"] = True
         raw_request["output_scores"] = True
@@ -112,15 +118,25 @@ class HuggingFaceClient(Client):
         self.cache = Cache(cache_config)
         self.model_server_instances: Dict[str, HuggingFaceServer] = {}
 
-    def get_model_server_instance(self, model_engine) -> HuggingFaceServer:
-        if model_engine not in self.model_server_instances:
-            if model_engine == "gpt-j-6b":
-                self.model_server_instances[model_engine] = HuggingFaceServer("EleutherAI/gpt-j-6B")
-            elif model_engine == "gpt2":
-                self.model_server_instances[model_engine] = HuggingFaceServer("gpt2")
+    def get_model_server_instance(self, model) -> HuggingFaceServer:
+        if model not in self.model_server_instances:
+            model_config = get_huggingface_model_config(model)
+            if model_config:
+                self.model_server_instances[model] = HuggingFaceServer(model_config)
+            elif model == "EleutherAI/gpt-j-6B":
+                self.model_server_instances[model] = HuggingFaceServer(
+                    HuggingFaceModelConfig.from_string("EleutherAI/gpt-j-6B")
+                )
+            elif model == "huggingface/gpt2":
+                self.model_server_instances[model] = HuggingFaceServer(HuggingFaceModelConfig.from_string("gpt2"))
+            elif model == "bigcode/santacoder":
+                self.model_server_instances[model] = HuggingFaceServer(
+                    HuggingFaceModelConfig.from_string("bigcode/santacoder")
+                )
             else:
-                raise Exception("Unknown model!")
-        return self.model_server_instances[model_engine]
+                raise Exception(f"Unknown HuggingFace model: {model}")
+
+        return self.model_server_instances[model]
 
     def make_request(self, request: Request) -> RequestResult:
         # Embedding not supported for this model
@@ -145,7 +161,7 @@ class HuggingFaceClient(Client):
 
         # Get cached model server instance if possible (to save on model and tokenizer
         # loading times).
-        model_server_instance: HuggingFaceServer = self.get_model_server_instance(request.model_engine)
+        model_server_instance: HuggingFaceServer = self.get_model_server_instance(request.model)
 
         try:
 

helm/proxy/clients/huggingface_model_registry.py

@@ -0,0 +1,111 @@
+from typing import Dict, Optional
+from dataclasses import dataclass
+import re
+from helm.common.hierarchical_logger import hlog
+from helm.proxy.models import (
+    Model,
+    ALL_MODELS,
+    MODEL_NAME_TO_MODEL,
+    TEXT_MODEL_TAG,
+    FULL_FUNCTIONALITY_TEXT_MODEL_TAG,
+)
+
+
+@dataclass(frozen=True)
+class HuggingFaceModelConfig:
+    namespace: Optional[str]
+    """Name of the group or user that owns the model. e.g. 'stanford-crfm'
+
+    May be None if the model (e.g. gpt2) does not have a namespace."""
+
+    model_name: str
+    """Name of the model. e.g. 'BioMedLM'
+
+    Does not include the namespace."""
+
+    revision: Optional[str]
+    """Revision of the model to use e.g. 'main'.
+
+    If None, use the default revision."""
+
+    @property
+    def model_id(self) -> str:
+        """Return the model ID.
+
+        Examples:
+        - 'gpt2'
+        - 'stanford-crfm/BioMedLM'"""
+        if self.namespace:
+            return f"{self.namespace}/{self.model_name}"
+        return self.model_name
+
+    def __str__(self) -> str:
+        """Return the full model name used by HELM in the format "[namespace/]model_name[@revision]".
+
+        Examples:
+        - 'gpt2'
+        - 'stanford-crfm/BioMedLM'
+        - 'stanford-crfm/BioMedLM@main'"""
+        result = self.model_name
+        if self.namespace:
+            result = f"{self.namespace}/{result}"
+        if self.revision:
+            result = f"{result}@{self.revision}"
+        return result
+
+    @staticmethod
+    def from_string(raw: str) -> "HuggingFaceModelConfig":
+        """Parses a string in the format "[namespace/]model_name[@revision]" to a HuggingFaceModelConfig.
+
+        Examples:
+        - 'gpt2'
+        - 'stanford-crfm/BioMedLM'
+        - 'stanford-crfm/BioMedLM@main'"""
+        pattern = r"((?P<namespace>[^/@]+)/)?(?P<model_name>[^/@]+)(@(?P<revision>[^/@]+))?"
+        match = re.fullmatch(pattern, raw)
+        if not match:
+            raise ValueError(f"Could not parse model name: '{raw}'; Expected format: [namespace/]model_name[@revision]")
+        model_name = match.group("model_name")
+        assert model_name
+        return HuggingFaceModelConfig(
+            namespace=match.group("namespace"), model_name=model_name, revision=match.group("revision")
+        )
+
+
+_huggingface_model_registry: Dict[str, HuggingFaceModelConfig] = {}
+
+
+def register_huggingface_model_config(model_name: str) -> HuggingFaceModelConfig:
+    """Register a AutoModelForCausalLM model from Hugging Face Model Hub for later use.
+
+    model_name format: namespace/model_name[@revision]"""
+    config = HuggingFaceModelConfig.from_string(model_name)
+    if config.model_id in _huggingface_model_registry:
+        raise ValueError(f"A Hugging Face model is already registered for model_id {model_name}")
+    _huggingface_model_registry[model_name] = config
+
+    # HELM model names require a namespace
+    if not config.namespace:
+        raise Exception("Registration of Hugging Face models without a namespace is not supported")
+    if model_name in MODEL_NAME_TO_MODEL:
+        raise ValueError(f"A HELM model is already registered for model name: {model_name}")
+    description = f"HuggingFace model {config.model_id}"
+    if config.revision:
+        description += f" at revision {config.revision}"
+    model = Model(
+        group=config.namespace,
+        name=model_name,
+        display_name=model_name,
+        creator_organization=config.namespace,
+        description=description,
+        tags=[TEXT_MODEL_TAG, FULL_FUNCTIONALITY_TEXT_MODEL_TAG],
+    )
+    MODEL_NAME_TO_MODEL[model_name] = model
+    ALL_MODELS.append(model)
+    hlog(f"Registered Hugging Face model: {model} config: {config}")
+    return config
+
+
+def get_huggingface_model_config(model_name: str) -> Optional[HuggingFaceModelConfig]:
+    """Returns a HuggingFaceModelConfig for the model_id."""
+    return _huggingface_model_registry.get(model_name)

helm/proxy/clients/huggingface_tokenizer.py

@@ -1,10 +1,12 @@
 import os
-from typing import Any, Dict
+from typing import Any, Dict, Optional
 
 from transformers import AutoTokenizer
 
 from helm.common.hierarchical_logger import htrack_block, hlog
 
+from helm.proxy.clients.huggingface_model_registry import get_huggingface_model_config
+
 
 class HuggingFaceTokenizers:
 
@@ -17,8 +19,11 @@ class HuggingFaceTokenizers:
         Returns the tokenizer.
         """
 
-        def load_tokenizer(hf_tokenizer_name: str):
+        def load_tokenizer(hf_tokenizer_name: str, revision: Optional[str] = None):
             """Loads tokenizer using files from disk if they exist. Otherwise, downloads from HuggingFace."""
+            tokenizer_kwargs = {}
+            if revision is not None:
+                tokenizer_kwargs["revision"] = revision
             try:
                 # From the Hugging Face documentation, "local_files_only(defaults to False) —
                 # Whether or not to only look at local files".
@@ -29,10 +34,14 @@ class HuggingFaceTokenizers:
                 # From https://huggingface.co/course/chapter6/3, "slow tokenizers are those written in Python inside
                 # the Hugging Face Transformers library, while the fast versions are the ones provided by Hugging Face
                 # Tokenizers, which are written in Rust." So, use the "fast" version of the tokenizers if available.
-                return AutoTokenizer.from_pretrained(hf_tokenizer_name, local_files_only=True, use_fast=True)
+                return AutoTokenizer.from_pretrained(
+                    hf_tokenizer_name, local_files_only=True, use_fast=True, **tokenizer_kwargs
+                )
             except OSError:
                 hlog(f"Local files do not exist for HuggingFace tokenizer: {hf_tokenizer_name}. Downloading...")
-                return AutoTokenizer.from_pretrained(hf_tokenizer_name, local_files_only=False, use_fast=True)
+                return AutoTokenizer.from_pretrained(
+                    hf_tokenizer_name, local_files_only=False, use_fast=True, **tokenizer_kwargs
+                )
 
         if tokenizer_name not in HuggingFaceTokenizers.tokenizers:
             with htrack_block(f"Loading {tokenizer_name} with Hugging Face Transformers"):
@@ -41,7 +50,12 @@ class HuggingFaceTokenizers:
 
                 # Weights are cached at ~/.cache/huggingface/transformers.
                 hf_tokenizer_name: str
-                if tokenizer_name == "huggingface/gpt2":
+                revision: Optional[str] = None
+                model_config = get_huggingface_model_config(tokenizer_name)
+                if model_config:
+                    hf_tokenizer_name = model_config.model_id
+                    revision = model_config.revision
+                elif tokenizer_name == "huggingface/gpt2":
                     hf_tokenizer_name = "gpt2"
                 elif tokenizer_name == "EleutherAI/gpt-j-6B":
                     # Not a typo: Named "gpt-j-6B" instead of "gpt-j-6b" in Hugging Face
@@ -58,10 +72,14 @@ class HuggingFaceTokenizers:
                     hf_tokenizer_name = "t5-11b"
                 elif tokenizer_name == "google/ul2":
                     hf_tokenizer_name = "google/ul2"
+                elif tokenizer_name == "google/flan-t5-xxl":
+                    hf_tokenizer_name = "google/flan-t5-xxl"
+                elif tokenizer_name == "bigcode/santacoder":
+                    hf_tokenizer_name = "bigcode/santacoder"
                 else:
-                    raise ValueError(f"Unsupported tokenizer: {tokenizer_name}")
+                    raise ValueError(f"Unsupported HuggingFace tokenizer: {tokenizer_name}")
 
                 # Keep the tokenizer in memory, so we don't recreate it for future requests
-                HuggingFaceTokenizers.tokenizers[tokenizer_name] = load_tokenizer(hf_tokenizer_name)
+                HuggingFaceTokenizers.tokenizers[tokenizer_name] = load_tokenizer(hf_tokenizer_name, revision)
 
         return HuggingFaceTokenizers.tokenizers[tokenizer_name]

helm/proxy/clients/openai_client.py

@@ -1,5 +1,5 @@
 from dataclasses import replace
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, cast
 
 import openai
 
@@ -24,6 +24,7 @@ class OpenAIClient(Client):
         self,
         api_key: str,
         cache_config: CacheConfig,
+        tokenizer_client: Client,
         chat_gpt_client: Optional[ChatGPTClient] = None,
         org_id: Optional[str] = None,
     ):
@@ -31,8 +32,12 @@ class OpenAIClient(Client):
         self.api_key: str = api_key
         self.api_base: str = "https://api.openai.com/v1"
         self.cache = Cache(cache_config)
+        self.tokenizer_client: Client = tokenizer_client
         self.chat_gpt_client: Optional[ChatGPTClient] = chat_gpt_client
 
+    def _is_chat_model_engine(self, model_engine: str):
+        return model_engine.startswith("gpt-3.5")
+
     def make_request(self, request: Request) -> RequestResult:
         if request.model_engine == "chat-gpt":
             assert self.chat_gpt_client is not None
@@ -44,6 +49,28 @@ class OpenAIClient(Client):
                 "input": request.prompt,
                 "engine": request.model_engine,
             }
+        elif self._is_chat_model_engine(request.model_engine):
+            raw_request = {
+                "model": request.model_engine,
+                # For now, put the whole prompt in a single user message, and expect the response
+                # to be returned in a single assistant message.
+                # TODO: Support ChatML for creating multiple messages with different roles.
+                # See: https://github.com/openai/openai-python/blob/main/chatml.md
+                "messages": [{"role": "user", "content": request.prompt}],
+                "temperature": request.temperature,
+                "top_p": request.top_p,
+                "n": request.num_completions,
+                # Note: Setting stop to ["\n"] results in an error
+                # See: https://community.openai.com/t/stop-n-in-gpt-3-5-turbo-leads-to-500-error/87815/15
+                # TODO: Handle this in the adapter.
+                "stop": request.stop_sequences or [],  # API doesn't like empty list
+                # Note: Chat models may require adding an extra token to max_tokens
+                # for the internal special role token.
+                # TODO: Handle this in the adapter.
+                "max_tokens": request.max_tokens,
+                "presence_penalty": request.presence_penalty,
+                "frequency_penalty": request.frequency_penalty,
+            }
         else:
             raw_request = {
                 "engine": request.model_engine,
@@ -74,6 +101,14 @@ class OpenAIClient(Client):
                     openai.api_base = self.api_base
                     return openai.Embedding.create(**raw_request)
 
+            elif self._is_chat_model_engine(request.model_engine):
+
+                def do_it():
+                    openai.organization = self.org_id
+                    openai.api_key = self.api_key
+                    openai.api_base = self.api_base
+                    return openai.ChatCompletion.create(**raw_request)
+
             else:
 
                 def do_it():
@@ -95,14 +130,37 @@ class OpenAIClient(Client):
         # needs to be populated, and `embedding` should be an empty list and vice-versa.
         embedding: List[float] = []
         completions: List[Sequence] = []
+        tokens: List[Token]
         if request.embedding:
             # If the user is requesting an embedding instead of completion
             # then completions would be left as an empty list. The embedding needs to be set.
             embedding = response["data"][0]["embedding"]
+        elif self._is_chat_model_engine(request.model_engine):
+            for raw_completion in response["choices"]:
+                # The ChatGPT API doesn't support echo. If `echo_prompt` is true, combine the prompt and completion.
+                raw_completion_content = raw_completion["message"]["content"]
+                text: str = request.prompt + raw_completion_content if request.echo_prompt else raw_completion_content
+                # The ChatGPT API doesn't return us tokens or logprobs, so we tokenize ourselves.
+                tokenization_result: TokenizationRequestResult = self.tokenizer_client.tokenize(
+                    # We're assuming ChatGPT uses the GPT-2 tokenizer.
+                    TokenizationRequest(text, tokenizer="huggingface/gpt2")
+                )
+                # Log probs are not currently not supported by the ChatGPT, so set to 0 for now.
+                tokens = [
+                    Token(text=cast(str, raw_token), logprob=0, top_logprobs={})
+                    for raw_token in tokenization_result.raw_tokens
+                ]
+                completion = Sequence(
+                    text=text,
+                    logprob=0,  # ChatGPT does not provide logprobs
+                    tokens=tokens,
+                    finish_reason={"reason": raw_completion["finish_reason"]},
+                )
+                completions.append(truncate_sequence(completion, request))  # Truncate the text by stop sequences
         else:
             for raw_completion in response["choices"]:
                 sequence_logprob = 0
-                tokens: List[Token] = []
+                tokens = []
 
                 raw_data = raw_completion["logprobs"]
                 for text, logprob, top_logprobs in zip(

helm/proxy/clients/test_huggingface_model_registry.py

@@ -0,0 +1,57 @@
+import pytest
+import unittest
+from typing import List, Tuple
+
+from helm.benchmark.run_expander import ModelRunExpander
+from helm.proxy.clients.huggingface_model_registry import (
+    HuggingFaceModelConfig,
+    register_huggingface_model_config,
+    get_huggingface_model_config,
+)
+from helm.proxy.models import get_all_models, get_all_text_models
+
+
+@pytest.mark.parametrize("model_name", ["EleutherAI/pythia-70m"])
+def test_hf_model_register(model_name):
+    register_huggingface_model_config(model_name)
+    assert model_name in ModelRunExpander("all").values
+    assert model_name in get_all_models()
+    assert model_name in get_all_text_models()
+
+
+class TestHuggingFaceModelRegistry(unittest.TestCase):
+    def test_round_trip(self):
+        config_pairs: List[Tuple[str, HuggingFaceModelConfig]] = [
+            ("gpt2", HuggingFaceModelConfig(namespace=None, model_name="gpt2", revision=None)),
+            (
+                "stanford-crfm/BioMedLM",
+                HuggingFaceModelConfig(namespace="stanford-crfm", model_name="BioMedLM", revision=None),
+            ),
+            (
+                "stanford-crfm/BioMedLM@main",
+                HuggingFaceModelConfig(namespace="stanford-crfm", model_name="BioMedLM", revision="main"),
+            ),
+        ]
+        for expected_model_name, expected_model_config in config_pairs:
+            actual_model_config = HuggingFaceModelConfig.from_string(expected_model_name)
+            actual_model_name = str(actual_model_config)
+            self.assertEqual(actual_model_name, expected_model_name)
+            self.assertEqual(actual_model_config, expected_model_config)
+
+    def test_model_id(self):
+        config_pairs: List[Tuple[str, str]] = [
+            ("gpt2", "gpt2"),
+            ("stanford-crfm/BioMedLM", "stanford-crfm/BioMedLM"),
+            ("stanford-crfm/BioMedLM@main", "stanford-crfm/BioMedLM"),
+        ]
+        for expected_model_name, expected_model_id in config_pairs:
+            actual_model_config = HuggingFaceModelConfig.from_string(expected_model_name)
+            self.assertEqual(actual_model_config.model_id, expected_model_id)
+
+    def test_register_huggingface_model_config(self):
+        register_huggingface_model_config("stanford-crfm/BioMedLM@main")
+        expected_model_config = HuggingFaceModelConfig(
+            namespace="stanford-crfm", model_name="BioMedLM", revision="main"
+        )
+        actual_model_config = get_huggingface_model_config("stanford-crfm/BioMedLM@main")
+        self.assertEqual(actual_model_config, expected_model_config)

helm/proxy/clients/test_huggingface_tokenizer.py

@@ -39,6 +39,9 @@ class TestHuggingFaceTokenizers:
     def test_get_tokenizer_ul2(self):
         TestHuggingFaceTokenizers.verify_get_tokenizer("google/ul2", 58)
 
+    def test_get_santacoder(self):
+        TestHuggingFaceTokenizers.verify_get_tokenizer("bigcode/santacoder", 62)
+
     def test_gpt2_tokenize_eos(self):
         eos_token: str = "<|endoftext|>"
         tokenizer = HuggingFaceTokenizers.get_tokenizer("huggingface/gpt2")

helm/proxy/clients/together_client.py

@@ -12,6 +12,22 @@ from helm.common.tokenization_request import (
 from .client import Client, wrap_request_time, truncate_sequence
 
 
+MODEL_ALIASES = {
+    "flan-t5-xxl": "flan-t5-xxl-hf",
+    "h3-2.7b": "h3-2.7b-h3",
+}
+"""Together model name aliases.
+
+HELM users use a shorter model name (e.g. together/flan-t5-xxl)
+whereas the Together client sends and caches requests using
+a longer model name that is suffixed with the implementation framework
+(e.g. flan-t5-xxl-hf). This allows trackcing exactly which
+implementation was used in the cached results, since some results may
+be different depending on the implementation (e.g. efficiency metrics).
+This also allows future migration of results in the case of changes of
+available implementations on Together."""
+
+
 def fix_text(x: str, model: str) -> str:
     """Fix text that comes back from the API."""
     x = x.replace("▁", " ")
@@ -31,8 +47,7 @@ class TogetherClient(Client):
         # Following the examples from https://github.com/togethercomputer/open-models-api
         return {
             "request_type": "language-model-inference",
-            # TODO: the Together API expects "together/" in the model name
-            "model": request.model,
+            "model": MODEL_ALIASES.get(request.model_engine, request.model_engine),
             "prompt": request.prompt,
             "temperature": request.temperature,
             "n": request.num_completions,

helm/proxy/clients/yalm_tokenizer/yalm_tokenizer.py

@@ -1,3 +1,4 @@
+import importlib_resources as resources
 import torch
 import sentencepiece as spm
 
@@ -7,6 +8,10 @@ adapted from https://github.com/yandex/YaLM-100B/blob/main/megatron_lm/megatron/
 """
 
 
+YALM_TOKENIZER_PACKAGE: str = "helm.proxy.clients.yalm_tokenizer"
+YALM_TOKENIZER_VOCAB_FILENAME: str = "voc_100b.sp"
+
+
 def convert_to_unicode(text):
     """Converts `text` to Unicode (if it's not already), assuming utf-8 input."""
     if isinstance(text, bytes):
@@ -28,9 +33,10 @@ class YaLMTokenizer:
     MASK_TOKEN = "[MASK]"
     MAX_SEQUENCE_LENGTH = 2048
 
-    def __init__(self, vocab_file="src/helm/proxy/clients/yalm_tokenizer/voc_100b.sp"):
+    def __init__(self):
         self.name = "sp"
-        self._tokenizer = spm.SentencePieceProcessor(model_file=vocab_file)
+        vocab_file_path = str(resources.files(YALM_TOKENIZER_PACKAGE).joinpath(YALM_TOKENIZER_VOCAB_FILENAME))
+        self._tokenizer = spm.SentencePieceProcessor(model_file=vocab_file_path)
         self._vocab_words = self._get_vocab_words()
         self.encoder = {token: idx for idx, token in enumerate(self._vocab_words)}
         self.decoder = {idx: token for idx, token in enumerate(self._vocab_words)}