crfm-helm 0.3.0__py3-none-any.whl → 0.5.0__py3-none-any.whl

helm/clients/mistral_client.py

@@ -0,0 +1,134 @@
+import requests
+from typing import Any, Dict, List, Optional, TypedDict
+
+from helm.proxy.retry import NonRetriableException
+from helm.common.cache import CacheConfig
+from helm.common.optional_dependencies import handle_module_not_found_error
+from helm.common.request import wrap_request_time, Request, RequestResult, GeneratedOutput
+from helm.tokenizers.tokenizer import Tokenizer
+from .client import CachingClient, truncate_and_tokenize_response_text
+
+try:
+    from mistralai.client import MistralClient
+    from mistralai.models.chat_completion import ChatMessage, ChatCompletionResponse
+except ModuleNotFoundError as e:
+    handle_module_not_found_error(e, ["mistral"])
+
+
+class MistralAIRequest(TypedDict):
+    """Data passed between make_request and _send_request. Used as the cache key."""
+
+    model: str
+    prompt: str
+    max_tokens: int
+    temperature: float
+    top_p: float
+    random_seed: Optional[int]
+
+
+class MistralAIClient(CachingClient):
+    """
+    Client for Mistral API.
+    """
+
+    def __init__(
+        self,
+        tokenizer: Tokenizer,
+        tokenizer_name: str,
+        cache_config: CacheConfig,
+        api_key: str,
+        mistral_model: Optional[str] = None,
+    ):
+        super().__init__(cache_config=cache_config)
+        self.api_key: str = api_key
+        self.tokenizer = tokenizer
+        self.tokenizer_name = tokenizer_name
+        self._client = MistralClient(api_key=self.api_key)
+        self.mistral_model = mistral_model
+
+    def _send_request(self, raw_request: MistralAIRequest) -> Dict[str, Any]:
+        messages = [ChatMessage(role="user", content=raw_request["prompt"])]
+
+        chat_response: ChatCompletionResponse = self._client.chat(
+            model=raw_request["model"],
+            messages=messages,
+            temperature=raw_request["temperature"],
+            max_tokens=raw_request["max_tokens"],
+            top_p=raw_request["top_p"],
+            random_seed=raw_request["random_seed"],
+            safe_prompt=False,  # Disable safe_prompt
+        )
+        # Documentation: "If mode is 'json', the output will only contain JSON serializable types."
+        # Source: https://docs.pydantic.dev/latest/api/base_model/#pydantic.BaseModel.model_dump
+        #
+        # We need to ensure that the output only contains JSON serializable types because the output
+        # will be serialized for storage in the cache.
+        return chat_response.model_dump(mode="json")
+
+    def _get_random_seed(self, request: Request, completion_index: int) -> Optional[int]:
+        if request.random is None and completion_index == 0:
+            return None
+
+        # Treat the user's request.random as an integer for the random seed.
+        try:
+            request_random_seed = int(request.random) if request.random is not None else 0
+        except ValueError:
+            raise NonRetriableException("MistralAIClient only supports integer values for request.random")
+
+        # A large prime is used so that the resulting values are unlikely to collide
+        # with request.random values chosen by the user.
+        fixed_large_prime = 1911011
+        completion_index_random_seed = completion_index * fixed_large_prime
+
+        return request_random_seed + completion_index_random_seed
+
+    def make_request(self, request: Request) -> RequestResult:
+        """Make a request"""
+        completions: List[GeneratedOutput] = []
+
+        # `num_completions` is not supported, so instead make `num_completions` separate requests.
+        for completion_index in range(request.num_completions):
+            try:
+                raw_request: MistralAIRequest = {
+                    "model": self.mistral_model or request.model_engine,
+                    "prompt": request.prompt,
+                    "max_tokens": request.max_tokens,
+                    "temperature": request.temperature,
+                    "top_p": request.top_p,
+                    "random_seed": self._get_random_seed(request, completion_index),
+                }
+
+                def do_it() -> Dict[str, Any]:
+                    result: Dict[str, Any] = self._send_request(raw_request)
+                    return result
+
+                # We need to include the engine's name to differentiate among requests made for different model
+                # engines since the engine name is not included in the request itself.
+                # In addition, we want to make `request.num_completions` fresh
+                # requests, cache key should contain the completion_index.
+                # Echoing the original prompt is not officially supported by Mistral. We instead prepend the
+                # completion with the prompt when `echo_prompt` is true, so keep track of it in the cache key.
+                cache_key = CachingClient.make_cache_key(raw_request, request)
+
+                response, cached = self.cache.get(cache_key, wrap_request_time(do_it))
+            except (requests.exceptions.RequestException, AssertionError) as e:
+                error: str = f"MistralClient error: {e}"
+                return RequestResult(success=False, cached=False, error=error, completions=[], embedding=[])
+
+            response_message: Dict[str, Any] = response["choices"][0]["message"]
+            assert response_message["role"] == "assistant"
+            response_text: str = response_message["content"]
+
+            # The Mistral API doesn't support echo. If `echo_prompt` is true, combine the prompt and completion.
+            text: str = request.prompt + response_text if request.echo_prompt else response_text
+            sequence = truncate_and_tokenize_response_text(text, request, self.tokenizer, self.tokenizer_name)
+            completions.append(sequence)
+
+        return RequestResult(
+            success=True,
+            cached=cached,
+            request_time=response["request_time"],
+            request_datetime=response["request_datetime"],
+            completions=completions,
+            embedding=[],
+        )

helm/clients/moderation_api_client.py

@@ -0,0 +1,109 @@
+from typing import Any, Dict
+
+from helm.common.request import wrap_request_time
+from helm.common.cache import Cache, CacheConfig
+from helm.common.moderations_api_request import (
+    ModerationCategoryScores,
+    ModerationCategoryFlaggedResults,
+    ModerationAPIRequest,
+    ModerationAPIRequestResult,
+)
+from helm.common.optional_dependencies import handle_module_not_found_error
+
+
+class ModerationAPIClient:
+    """
+    From https://beta.openai.com/docs/guides/moderation/overview, the moderation endpoint is a tool
+    to check whether content complies with OpenAI's content policy. Developers can thus identify content
+    that OpenAI's content policy prohibits and take action, for instance by filtering it.
+    """
+
+    # For descriptions of the models, see https://beta.openai.com/docs/api-reference/moderations/create
+    LATEST_MODEL: str = "text-moderation-latest"
+    STABLE_MODEL: str = "text-moderation-stable"
+
+    # List of categories (https://beta.openai.com/docs/guides/moderation/overview)
+    HATE: str = "hate"
+    HATE_THREATENING: str = "hate/threatening"
+    SELF_HARM: str = "self-harm"
+    SEXUAL: str = "sexual"
+    SEXUAL_MINORS: str = "sexual/minors"
+    VIOLENCE: str = "violence"
+    VIOLENCE_GRAPHIC: str = "violence/graphic"
+
+    def __init__(self, api_key: str, cache_config: CacheConfig):
+        self.cache = Cache(cache_config)
+        try:
+            from openai import OpenAI
+        except ModuleNotFoundError as e:
+            handle_module_not_found_error(e, ["openai"])
+        # TODO: Add OpenAI organization.
+        self.client = OpenAI(api_key=api_key)
+
+    def get_moderation_results(self, request: ModerationAPIRequest) -> ModerationAPIRequestResult:
+        """
+        Sends a request to OpenAI's moderation endpoint.
+        https://beta.openai.com/docs/api-reference/moderations/create
+        """
+        try:
+            import openai
+        except ModuleNotFoundError as e:
+            handle_module_not_found_error(e, ["openai"])
+
+        raw_request: Dict[str, str] = {
+            "input": request.text,
+            "model": self.LATEST_MODEL if request.use_latest_model else self.STABLE_MODEL,
+        }
+
+        try:
+
+            def do_it() -> Dict[str, Any]:
+                result = self.client.moderations.create(input=request.text).model_dump(mode="json")
+                assert "results" in result and len(result["results"]) > 0, f"Invalid response: {result}"
+                return result
+
+            response, cached = self.cache.get(raw_request, wrap_request_time(do_it))
+        except openai.OpenAIError as e:
+            error: str = f"Moderation API error: {e}"
+            return ModerationAPIRequestResult(
+                success=False, cached=False, error=error, flagged=None, flagged_results=None, scores=None
+            )
+
+        moderation_result = response["results"][0]
+        category_results: Dict[str, bool] = moderation_result["categories"]
+        score_results: Dict[str, float] = moderation_result["category_scores"]
+
+        flagged_results = ModerationCategoryFlaggedResults(
+            hate_flagged=category_results[self.HATE],
+            hate_threatening_flagged=category_results[self.HATE_THREATENING],
+            self_harm_flagged=category_results[self.SELF_HARM],
+            sexual_flagged=category_results[self.SEXUAL],
+            sexual_minors_flagged=category_results[self.SEXUAL_MINORS],
+            violence_flagged=category_results[self.VIOLENCE],
+            violence_graphic_flagged=category_results[self.VIOLENCE_GRAPHIC],
+        )
+        scores = ModerationCategoryScores(
+            hate_score=score_results[self.HATE],
+            hate_threatening_score=score_results[self.HATE_THREATENING],
+            self_harm_score=score_results[self.SELF_HARM],
+            sexual_score=score_results[self.SEXUAL],
+            sexual_minors_score=score_results[self.SEXUAL_MINORS],
+            violence_score=score_results[self.VIOLENCE],
+            violence_graphic_score=score_results[self.VIOLENCE_GRAPHIC],
+        )
+        return ModerationAPIRequestResult(
+            success=True,
+            cached=cached,
+            flagged=moderation_result["flagged"],
+            flagged_results=flagged_results,
+            scores=scores,
+        )
+
+    def will_be_flagged(self, text: str) -> bool:
+        """Returns True if the text is against OpenAI's content policy and will be flagged, False otherwise."""
+        result: ModerationAPIRequestResult = self.get_moderation_results(
+            # Use the latest model so the account does not get banned
+            ModerationAPIRequest(text=text, use_latest_model=True)
+        )
+        assert result.flagged is not None
+        return result.flagged

helm/clients/open_lm_client.py

@@ -0,0 +1,43 @@
+from threading import Lock
+from typing import Optional
+
+from transformers import AutoConfig, AutoModelForCausalLM
+from helm.common.cache import CacheConfig
+
+from helm.common.optional_dependencies import OptionalDependencyNotInstalled
+from helm.clients.huggingface_client import HuggingFaceClient
+
+
+_register_open_lm_lock = Lock()
+_register_open_lm_done = False
+
+
+def _register_open_lm_for_auto_model():
+    """Register OpenLMForCausalLM for AutoModelForCausalLM."""
+    try:
+        from open_lm.utils.transformers.hf_model import OpenLMforCausalLM
+        from open_lm.utils.transformers.hf_config import OpenLMConfig
+    except ModuleNotFoundError as e:
+        # Provide manual instructions for installing open_lm from GitHub
+        # because PyPI does not allow installing dependencies directly from GitHub.
+        raise OptionalDependencyNotInstalled(
+            f"Optional dependency {e.name} is not installed. "
+            "Please run `pip install open_lm@git+https://github.com/mlfoundations/open_lm.git@main` to install it."
+        ) from e
+
+    with _register_open_lm_lock:
+        global _register_open_lm_done
+        if not _register_open_lm_done:
+            AutoConfig.register("openlm", OpenLMConfig)
+            AutoModelForCausalLM.register(OpenLMConfig, OpenLMforCausalLM)
+        _register_open_lm_done = True
+
+
+class OpenLMClient(HuggingFaceClient):
+    """Client for OpenLM: https://github.com/mlfoundations/open_lm"""
+
+    def __init__(self, cache_config: CacheConfig, pretrained_model_name_or_path: Optional[str] = None, **kwargs):
+        _register_open_lm_for_auto_model()
+        super().__init__(
+            cache_config=cache_config, pretrained_model_name_or_path=pretrained_model_name_or_path, **kwargs
+        )

helm/clients/openai_client.py

@@ -0,0 +1,302 @@
+# mypy: check_untyped_defs = False
+from dataclasses import replace
+from typing import Any, Dict, List, Optional, cast, Union
+
+from helm.benchmark.model_metadata_registry import is_vlm
+from helm.common.cache import CacheConfig
+from helm.common.media_object import TEXT_TYPE
+from helm.common.request import wrap_request_time, Request, RequestResult, GeneratedOutput, Token
+from helm.common.hierarchical_logger import hlog
+from helm.common.optional_dependencies import handle_module_not_found_error
+from helm.common.tokenization_request import (
+    TokenizationRequest,
+    TokenizationRequestResult,
+)
+from helm.tokenizers.tokenizer import Tokenizer
+from .client import CachingClient, truncate_sequence, generate_uid_for_multimodal_prompt
+
+try:
+    import openai
+    from openai import OpenAI
+except ModuleNotFoundError as e:
+    handle_module_not_found_error(e, ["openai"])
+
+
+class OpenAIClient(CachingClient):
+    END_OF_TEXT: str = "<|endoftext|>"
+
+    # Error OpenAI throws when the image in the prompt violates their content policy
+    INAPPROPRIATE_IMAGE_ERROR: str = "Your input image may contain content that is not allowed by our safety system"
+
+    # Set the finish reason to this if the prompt violates OpenAI's content policy
+    CONTENT_POLICY_VIOLATED_FINISH_REASON: str = (
+        "The prompt violates OpenAI's content policy. "
+        "See https://labs.openai.com/policies/content-policy for more information."
+    )
+
+    def __init__(
+        self,
+        tokenizer: Tokenizer,
+        tokenizer_name: str,
+        cache_config: CacheConfig,
+        api_key: Optional[str] = None,
+        org_id: Optional[str] = None,
+        base_url: Optional[str] = None,
+    ):
+        super().__init__(cache_config=cache_config)
+        self.tokenizer = tokenizer
+        self.tokenizer_name = tokenizer_name
+        self.client = OpenAI(api_key=api_key, organization=org_id, base_url=base_url)
+
+    def _is_chat_model_engine(self, model_engine: str) -> bool:
+        if model_engine == "gpt-3.5-turbo-instruct":
+            return False
+        elif model_engine.startswith("gpt-3.5") or model_engine.startswith("gpt-4"):
+            return True
+        return False
+
+    def _get_model_for_request(self, request: Request) -> str:
+        return request.model_engine
+
+    def _get_cache_key(self, raw_request: Dict, request: Request):
+        cache_key = CachingClient.make_cache_key(raw_request, request)
+        if is_vlm(request.model):
+            assert request.multimodal_prompt is not None
+            prompt_key: str = generate_uid_for_multimodal_prompt(request.multimodal_prompt)
+            cache_key = {**cache_key, "multimodal_prompt": prompt_key}
+            del cache_key["messages"]
+        return cache_key
+
+    def _make_embedding_request(self, request: Request) -> RequestResult:
+        raw_request: Dict[str, Any]
+        raw_request = {
+            "input": request.prompt,
+            # Note: In older deprecated versions of the OpenAI API, "model" used to be "engine".
+            "model": self._get_model_for_request(request),
+        }
+
+        def do_it() -> Dict[str, Any]:
+            return self.client.embeddings.create(**raw_request).model_dump(mode="json")
+
+        try:
+            cache_key = self._get_cache_key(raw_request, request)
+            response, cached = self.cache.get(cache_key, wrap_request_time(do_it))
+        except openai.OpenAIError as e:
+            error: str = f"OpenAI error: {e}"
+            return RequestResult(success=False, cached=False, error=error, completions=[], embedding=[])
+
+        # If the user is requesting completions instead of an embedding, then `completions`
+        # needs to be populated, and `embedding` should be an empty list and vice-versa.
+        embedding: List[float] = []
+        # 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"]
+
+        return RequestResult(
+            success=True,
+            cached=cached,
+            request_time=response["request_time"],
+            request_datetime=response.get("request_datetime"),
+            completions=[],
+            embedding=embedding,
+        )
+
+    def _make_chat_request(self, request: Request) -> RequestResult:
+        messages: Optional[List[Dict[str, Union[str, Any]]]] = request.messages
+        if request.messages is not None:
+            # Checks that all messages have a role and some content
+            for message in request.messages:
+                if not message.get("role") or not message.get("content"):
+                    raise ValueError("All messages must have a role and content")
+            # Checks that the last role is "user"
+            if request.messages[-1]["role"] != "user":
+                raise ValueError("Last message must have role 'user'")
+            if request.prompt != "":
+                hlog("WARNING: Since message is set, prompt will be ignored")
+        else:
+            # Convert prompt into a single message
+            # 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
+
+            # Content can either be text or a list of multimodal content made up of text and images:
+            # https://platform.openai.com/docs/guides/vision
+            content: Union[str, List[Union[str, Any]]]
+            if request.multimodal_prompt is not None:
+                content = []
+                for media_object in request.multimodal_prompt.media_objects:
+                    if media_object.is_type("image") and media_object.location:
+                        from helm.common.images_utils import encode_base64
+
+                        base64_image: str = encode_base64(media_object.location)
+                        content.append(
+                            {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"}}
+                        )
+                    elif media_object.is_type(TEXT_TYPE):
+                        if media_object.text is None:
+                            raise ValueError("MediaObject of text type has missing text field value")
+                        content.append({"type": media_object.type, "text": media_object.text})
+                    else:
+                        raise ValueError(f"Unrecognized MediaObject type {media_object.type}")
+
+            else:
+                content = request.prompt
+
+            messages = [{"role": "user", "content": content}]
+
+        raw_request: Dict[str, Any] = {
+            "model": self._get_model_for_request(request),
+            "messages": messages,
+            "temperature": request.temperature,
+            "top_p": request.top_p,
+            "n": request.num_completions,
+            "stop": request.stop_sequences or None,  # API doesn't like empty list
+            # Note: Chat models may require adding an extra token to max_tokens
+            # for the internal special role token.
+            "max_tokens": request.max_tokens,
+            "presence_penalty": request.presence_penalty,
+            "frequency_penalty": request.frequency_penalty,
+        }
+
+        # OpenAI's vision API doesn't allow None values for stop.
+        # Fails with "body -> stop: none is not an allowed value" if None is passed.
+        if is_vlm(request.model) and raw_request["stop"] is None:
+            raw_request.pop("stop")
+
+        def do_it() -> Dict[str, Any]:
+            return self.client.chat.completions.create(**raw_request).model_dump(mode="json")
+
+        try:
+            cache_key = self._get_cache_key(raw_request, request)
+            response, cached = self.cache.get(cache_key, wrap_request_time(do_it))
+        except openai.OpenAIError as e:
+            if self.INAPPROPRIATE_IMAGE_ERROR in str(e):
+                hlog(f"Failed safety check: {str(request)}")
+                empty_completion = GeneratedOutput(
+                    text="",
+                    logprob=0,
+                    tokens=[],
+                    finish_reason={"reason": self.CONTENT_POLICY_VIOLATED_FINISH_REASON},
+                )
+                return RequestResult(
+                    success=True,
+                    cached=False,
+                    request_time=0,
+                    completions=[empty_completion] * request.num_completions,
+                    embedding=[],
+                )
+
+            error: str = f"OpenAI error: {e}"
+            return RequestResult(success=False, cached=False, error=error, completions=[], embedding=[])
+
+        completions: List[GeneratedOutput] = []
+        for raw_completion in response["choices"]:
+            # The OpenAI chat completion 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 OpenAI chat completion API doesn't return us tokens or logprobs, so we tokenize ourselves.
+            tokenization_result: TokenizationRequestResult = self.tokenizer.tokenize(
+                TokenizationRequest(text, tokenizer=self.tokenizer_name)
+            )
+            # Log probs are not currently not supported by the OpenAI chat completion API, so set to 0 for now.
+            tokens: List[Token] = [
+                Token(text=cast(str, raw_token), logprob=0) for raw_token in tokenization_result.raw_tokens
+            ]
+            completion = GeneratedOutput(
+                text=text,
+                logprob=0,  # OpenAI 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
+
+        return RequestResult(
+            success=True,
+            cached=cached,
+            request_time=response["request_time"],
+            request_datetime=response.get("request_datetime"),
+            completions=completions,
+            embedding=[],
+        )
+
+    def _to_raw_completion_request(self, request: Request) -> Dict[str, Any]:
+        raw_request: Dict[str, Any] = {
+            # Note: In older deprecated versions of the OpenAI API, "model" used to be "engine".
+            "model": self._get_model_for_request(request),
+            "prompt": request.prompt,
+            "temperature": request.temperature,
+            "n": request.num_completions,
+            "max_tokens": request.max_tokens,
+            "best_of": request.top_k_per_token,
+            "logprobs": request.top_k_per_token,
+            "stop": request.stop_sequences or None,  # API doesn't like empty list
+            "top_p": request.top_p,
+            "presence_penalty": request.presence_penalty,
+            "frequency_penalty": request.frequency_penalty,
+            "echo": request.echo_prompt,
+        }
+
+        # OpenAI doesn't let you ask for more completions than the number of
+        # per-token candidates.
+        raw_request["best_of"] = max(raw_request["best_of"], raw_request["n"])
+        raw_request["logprobs"] = max(raw_request["logprobs"], raw_request["n"])
+
+        return raw_request
+
+    def _make_completion_request(self, request: Request) -> RequestResult:
+        raw_request = self._to_raw_completion_request(request)
+
+        def do_it() -> Dict[str, Any]:
+            return self.client.completions.create(**raw_request).model_dump(mode="json")
+
+        try:
+            cache_key = self._get_cache_key(raw_request, request)
+            response, cached = self.cache.get(cache_key, wrap_request_time(do_it))
+        except openai.OpenAIError as e:
+            error: str = f"OpenAI error: {e}"
+            return RequestResult(success=False, cached=False, error=error, completions=[], embedding=[])
+
+        completions: List[GeneratedOutput] = []
+        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))
+                sequence_logprob += logprob or 0
+            completion = GeneratedOutput(
+                text=raw_completion["text"],
+                logprob=sequence_logprob,
+                tokens=tokens,
+                finish_reason={"reason": raw_completion["finish_reason"]},
+            )
+            # OpenAI sends us back tokens past the end of text token,
+            # so we need to manually truncate the list of tokens.
+            # TODO: filed an issue with their support to check what the expected behavior here is.
+            completion = truncate_sequence(
+                completion, replace(request, stop_sequences=request.stop_sequences + [OpenAIClient.END_OF_TEXT])
+            )
+            completions.append(completion)
+
+        return RequestResult(
+            success=True,
+            cached=cached,
+            request_time=response["request_time"],
+            request_datetime=response.get("request_datetime"),
+            completions=completions,
+            embedding=[],
+        )
+
+    def make_request(self, request: Request) -> RequestResult:
+        if request.embedding:
+            return self._make_embedding_request(request)
+        elif self._is_chat_model_engine(request.model_engine):
+            return self._make_chat_request(request)
+        else:
+            return self._make_completion_request(request)

helm/{proxy/clients → clients}/palmyra_client.py

@@ -5,12 +5,12 @@ from typing import Any, Dict, List
 
 from helm.common.cache import CacheConfig
 from helm.common.hierarchical_logger import hlog
-from helm.common.request import wrap_request_time, Request, RequestResult, Sequence, Token, ErrorFlags
+from helm.common.request import wrap_request_time, Request, RequestResult, GeneratedOutput, Token, ErrorFlags
 from helm.common.tokenization_request import (
     TokenizationRequest,
     TokenizationRequestResult,
 )
-from helm.proxy.tokenizers.tokenizer import Tokenizer
+from helm.tokenizers.tokenizer import Tokenizer
 from .client import CachingClient, truncate_sequence
 
 
@@ -28,9 +28,11 @@ def _is_content_moderation_failure(response: Dict) -> bool:
 
 
 class PalmyraClient(CachingClient):
-    def __init__(self, api_key: str, tokenizer: Tokenizer, cache_config: CacheConfig):
-        super().__init__(cache_config=cache_config, tokenizer=tokenizer)
+    def __init__(self, tokenizer: Tokenizer, tokenizer_name: str, cache_config: CacheConfig, api_key: str):
+        super().__init__(cache_config=cache_config)
         self.api_key: str = api_key
+        self.tokenizer = tokenizer
+        self.tokenizer_name = tokenizer_name
 
     def _send_request(self, model_name: str, raw_request: Dict[str, Any]) -> Dict[str, Any]:
         response = requests.request(
@@ -65,14 +67,14 @@ class PalmyraClient(CachingClient):
             # "random_seed": request.random,
         }
 
-        completions: List[Sequence] = []
+        completions: List[GeneratedOutput] = []
         model_name: str = request.model_engine
 
         # `num_completions` is not supported, so instead make `num_completions` separate requests.
         for completion_index in range(request.num_completions):
             try:
 
-                def do_it():
+                def do_it() -> Dict[str, Any]:
                     # Add an argument timeout to raw_request to avoid waiting getting timeout of 60s
                     # which happens for long prompts.
                     request_with_timeout = {"timeout": 300, **raw_request}
@@ -100,7 +102,10 @@ class PalmyraClient(CachingClient):
                 return RequestResult(success=False, cached=False, error=error, completions=[], embedding=[])
 
             if _is_content_moderation_failure(response):
-                hlog(f"WARNING: Returning empty request for {request.model} due to content moderation filter")
+                hlog(
+                    f"WARNING: Returning empty request for {request.model_deployment} "
+                    "due to content moderation filter"
+                )
                 return RequestResult(
                     success=False,
                     cached=False,
@@ -119,15 +124,13 @@ class PalmyraClient(CachingClient):
             # The Writer API doesn't return us tokens or logprobs, so we tokenize ourselves.
             tokenization_result: TokenizationRequestResult = self.tokenizer.tokenize(
                 # Writer uses the GPT-2 tokenizer
-                TokenizationRequest(text, tokenizer="huggingface/gpt2")
+                TokenizationRequest(text, tokenizer=self.tokenizer_name)
             )
 
             # Log probs are not currently not supported by the Writer, so set to 0 for now.
-            tokens: List[Token] = [
-                Token(text=str(text), logprob=0, top_logprobs={}) for text in tokenization_result.raw_tokens
-            ]
+            tokens: List[Token] = [Token(text=str(text), logprob=0) for text in tokenization_result.raw_tokens]
 
-            completion = Sequence(text=response_text, logprob=0, tokens=tokens)
+            completion = GeneratedOutput(text=response_text, logprob=0, tokens=tokens)
             sequence = truncate_sequence(completion, request, print_warning=True)
             completions.append(sequence)