Repository: stanford-crfm/helm
Issue: Support local tokenizer for Aleph Alpha
Root cause: Unimplemented Feature Gap

Integrate into our codebase (also notice the `nice` parameter):

```python3
import json
import requests
import time
from tokenizers import Tokenizer
from typing import Any, Dict, List

from helm.common.cache import Cache, CacheConfig
from helm.common.request import Request, RequestResult, Sequence, Token
from helm.common.tokenization_request import (
    DecodeRequest,
    DecodeRequestResult,
    TokenizationRequest,
    TokenizationRequestResult,
    TokenizationToken,
)
from .client import Client, wrap_request_time, truncate_sequence


class AlephAlphaClient(Client):
    COMPLETION_ENDPOINT: str = "complete"
    TOKENIZE_ENDPOINT: str = "tokenize"
    DETOKENIZE_ENDPOINT: str = "detokenize"

    def __init__(self, api_key: str, cache_config: CacheConfig):
        self.api_key: str = api_key
        self.cache = Cache(cache_config)
        # instantiate standalone tokenizer
        self.tokenizer = Tokenizer.from_file("alpha-001-128k.json")
        self.num_requests = 0  # counter for debugging

    def _send_request(self, endpoint: str, raw_request: Dict[str, Any]) -> Dict[str, Any]:
        num_trials = 0
        max_trials = 100
        sleep_time = 60
        while num_trials < max_trials:
            num_trials += 1
            self.num_requests += 1
            #print("Number of requests so far:", self.num_requests, flush=True)

            try:
                params = {
                    "nice": "true",
                }
                response = requests.request(
                    method="POST",
                    url=f"https://api.aleph-alpha.com/{endpoint}",
                    headers={
                        "Content-Type": "application/json",
                        "Accept": "application/json",
                        "Authorization": f"Bearer {self.api_key}",
                    },
                    params=params,
                    data=json.dumps(raw_request),
                )
                response.raise_for_status()  # this will raise an exception if the response was not OK and therefore the completion didn't go through
                result = json.loads(response.text)
            
            except TimeoutError as e:
                print(f"Client TimeoutError; sleep for {sleep_time} seconds", flush=True)
                time.sleep(sleep_time)
                sleep_time *= 1.5

            except RuntimeError as e:
                print(f"Client RuntimeError; sleep for {sleep_time} seconds", flush=True)
                time.sleep(sleep_time)
                sleep_time *= 1.5

            except Exception as e:
                print(response.text, flush=True)
                raise

        assert "error" not in result, f"Request failed with error: {result['error']}"
        return result

    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.

        if request.max_tokens == 0:
            print("Error: max_tokens = 0", flush=True)
        #assert request.max_tokens > 0, "only max_tokens > 0 is supported by the Aleph Alpha API"

        raw_request = {
            "model": request.model_engine,
            "prompt": request.prompt,
            "maximum_tokens": request.max_tokens,
            "temperature": request.temperature,
            "top_k": request.top_k_per_token,
            "top_p": request.top_p,
            "presence_penalty": request.presence_penalty,
            "frequency_penalty": request.frequency_penalty,
            "n": request.num_completions,
            "stop_sequences": request.stop_sequences,
            "log_probs": request.top_k_per_token,
            "tokens": True,  # Setting to True returns individual tokens of the completion
        }

        try:
            def do_it():
                result = self._send_request(AlephAlphaClient.COMPLETION_ENDPOINT, raw_request)
                assert "completions" 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}"
            return RequestResult(success=False, cached=False, error=error, completions=[], embedding=[])

        completions: List[Sequence] = []
        for completion in response["completions"]:
            sequence_logprob: float = 0
            tokens: List[Token] = []

            # `completion_tokens` is the list of selected tokens.
            for i, token in enumerate(completion["completion_tokens"]):
                # Get the top K logprobs for the ith token
                top_logprobs: Dict[str, float] = completion["log_probs"][i]
                # Use the selected token value to get the logprob
                logprob: float = top_logprobs[token]
                sequence_logprob += logprob
                tokens.append(
                    Token(
                        text=token,
                        logprob=logprob,
                        top_logprobs=top_logprobs,
                    )
                )

            sequence: Sequence = Sequence(text=completion["completion"], logprob=sequence_logprob, tokens=tokens)
            sequence = truncate_sequence(sequence, request)
            completions.append(sequence)

        return RequestResult(
            success=True,
            cached=cached,
            request_time=response["request_time"],
            request_datetime=response["request_datetime"],
            completions=completions,
            embedding=[],
        )

    def tokenize(self, request: TokenizationRequest) -> TokenizationRequestResult:
        """Tokenize a string with the standalone tokenizer."""
        raw_request = {
            "model": request.tokenizer_name,
            "prompt": request.text,
            "tokens": True,
            "token_ids": True,
        }
        try:

            def do_it():
                encoded = self.tokenizer.encode(raw_request["prompt"], add_special_tokens=False)
                result = {
                    "token_ids": encoded.ids,
                    "tokens": encoded.tokens,
                }
                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}"
            return TokenizationRequestResult(error=error, success=False, cached=False, text="", tokens=[])

        tokens = response["token_ids" if request.encode else "tokens"]

        if request.truncation:
            tokens = tokens[: request.max_length]

        return TokenizationRequestResult(
            success=True,
            cached=cached,
            tokens=[TokenizationToken(value) for value in tokens],
            text=request.text,
            request_time=response["request_time"],
        )

    def tokenize_api(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,
        }

        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))
            pass
        except (requests.exceptions.RequestException, AssertionError) as e:
            error: str = f"AlephAlphaClient error: {e}"
            return TokenizationRequestResult(error=error, success=False, cached=False, text="", tokens=[])

        tokens = response["token_ids" if request.encode else "tokens"]
        if request.truncation:
            tokens = tokens[: request.max_length]

        return TokenizationRequestResult(
            success=True,
            cached=cached,
            tokens=[TokenizationToken(value) for value in tokens],
            text=request.text,
            request_time=response["request_time"],
        )

    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,
        }

        try:

            def do_it():
                result = {
                    "result": self.tokenizer.decode(request.tokens, skip_special_tokens=True),
                }
                assert "result" 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}"
            return DecodeRequestResult(error=error, success=False, cached=False, text="")

        return DecodeRequestResult(
            success=True,
            cached=cached,
            # The text always seems to start with a single whitespace when encoding/decoding.
            text=response["result"].replace(" ", "", 1),
            request_time=response["request_time"],
        )

    def decode_api(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,
        }

        try:

            def do_it():
                result = self._send_request(AlephAlphaClient.DETOKENIZE_ENDPOINT, raw_request)
                assert "result" 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}"
            return DecodeRequestResult(error=error, success=False, cached=False, text="")

        return DecodeRequestResult(
            success=True,
            cached=cached,
            # The text always seems to start with a single whitespace when encoding/decoding.
            text=response["result"].replace(" ", "", 1),
            request_time=response["request_time"],
        )

```

Aleph Alpha just released their tokenizer as part of their Python package:

https://aleph-alpha-client.readthedocs.io/en/latest/aleph_alpha_client.html#aleph_alpha_client.Client.tokenizer