crfm-helm 0.4.0__py3-none-any.whl → 0.5.10__py3-none-any.whl

helm/clients/bedrock_client.py

@@ -0,0 +1,381 @@
+from abc import abstractmethod
+from copy import deepcopy
+import json
+import os
+from typing import Any, Dict, List, Mapping, Optional, TypedDict
+from datetime import datetime
+
+from helm.common.cache import CacheConfig
+from helm.clients.client import CachingClient, truncate_and_tokenize_response_text
+from helm.common.hierarchical_logger import hexception
+from helm.common.request import Request, RequestResult, GeneratedOutput, wrap_request_time
+from helm.clients.bedrock_utils import get_bedrock_client, get_bedrock_client_v1
+from helm.proxy.retry import NonRetriableException
+from helm.tokenizers.tokenizer import Tokenizer
+
+
+JSON_CONTENT_TYPE = "application/json"
+
+
+class BedrockClient(CachingClient):
+    @abstractmethod
+    def convert_request_to_raw_request(self, request: Request) -> Dict:
+        raise NotImplementedError()
+
+    @abstractmethod
+    def convert_raw_response_to_completions(self, response: Dict, request: Request) -> List[GeneratedOutput]:
+        raise NotImplementedError()
+
+    """
+    Amazon Bedrock is a fully managed service that provides s selection of leading foundation models (FMs) from Amazon
+    and other partner model providers.
+    """
+
+    @property
+    @abstractmethod
+    def model_provider(self) -> str:
+        raise NotImplementedError()
+
+    def __init__(
+        self,
+        cache_config: CacheConfig,
+        tokenizer: Tokenizer,
+        tokenizer_name: str,
+        assumed_role: Optional[str] = None,
+        region: Optional[str] = None,
+        bedrock_model_id: Optional[str] = None,
+    ):
+        super().__init__(cache_config=cache_config)
+        self.tokenizer = tokenizer
+        self.tokenizer_name = tokenizer_name
+        self.bedrock_client = get_bedrock_client(
+            assumed_role=assumed_role or os.environ.get("BEDROCK_ASSUME_ROLE", None),
+            region=region,
+        )
+        self.bedrock_model_id = bedrock_model_id
+
+    def make_request(self, request: Request) -> RequestResult:
+        if self.bedrock_model_id:
+            model_id = self.bedrock_model_id
+        else:
+            # model_id should be something like "amazon.titan-tg1-large", replace amazon- prefix with model creator name
+            model_name = request.model.split("/")[-1]
+            # check if model_name starts with "amazon-"
+            if self.model_provider == "amazon":
+                model_id = f"{self.model_provider}.{model_name}"
+            else:
+                model_id = model_name.replace("amazon-", f"{self.model_provider}.")
+
+        raw_request = self.convert_request_to_raw_request(request)
+
+        # modelId isn't part of raw_request, so it must be explicitly passed into the input to
+        raw_request_for_cache: Dict = {"modelId": model_id, **deepcopy(raw_request)}
+        cache_key: Mapping = CachingClient.make_cache_key(raw_request_for_cache, request)
+
+        def do_it() -> Dict[Any, Any]:
+            response = self.bedrock_client.invoke_model(
+                body=json.dumps(raw_request), modelId=model_id, accept=JSON_CONTENT_TYPE, contentType=JSON_CONTENT_TYPE
+            )
+            return json.loads(response.get("body").read())
+
+        try:
+            response, cached = self.cache.get(cache_key, wrap_request_time(do_it))
+
+        except Exception as error:
+            hexception(error)
+            return RequestResult(
+                success=False,
+                cached=False,
+                error=str(error),
+                completions=[],
+                embedding=[],
+            )
+
+        completions = self.convert_raw_response_to_completions(response, request)
+
+        return RequestResult(
+            success=True,
+            cached=cached,
+            request_time=response["request_time"],
+            request_datetime=response["request_datetime"],
+            completions=completions,
+            embedding=[],
+        )
+
+
+class _ContentBlock(TypedDict):
+    text: str
+
+
+class _Message(TypedDict):
+    role: str
+    content: List[_ContentBlock]
+
+
+class BedrockNovaClient(CachingClient):
+    """
+    Amazon Bedrock is a fully managed service that provides s selection of leading foundation models (FMs) from Amazon
+    and other partner model providers.
+    """
+
+    def __init__(
+        self,
+        cache_config: CacheConfig,
+        tokenizer: Tokenizer,
+        tokenizer_name: str,
+        assumed_role: Optional[str] = None,
+        region: Optional[str] = None,
+        bedrock_model_id: Optional[str] = None,
+    ):
+        super().__init__(cache_config=cache_config)
+        self.tokenizer = tokenizer
+        self.tokenizer_name = tokenizer_name
+        self.bedrock_model_id = bedrock_model_id
+        self.bedrock_client = get_bedrock_client_v1(
+            assumed_role=assumed_role or os.environ.get("BEDROCK_ASSUME_ROLE", None),
+            region=region,
+        )
+
+    def _get_messages_from_request(self, request: Request) -> List[_Message]:
+        if request.prompt and request.messages:
+            raise ValueError(f"Only one of `prompt` and `messages` may be set in request: {request}")
+        if request.multimodal_prompt:
+            raise ValueError(f"`multimodal_prompt` is not supported in request: {request}")
+
+        if request.messages:
+            return [
+                {"role": message["role"], "content": [{"text": message["content"]}]} for message in request.messages
+            ]
+        else:
+            return [{"role": "user", "content": [{"text": request.prompt}]}]
+
+    def convert_request_to_raw_request(self, request: Request) -> Dict:
+        model_id = request.model.replace("/", ".")
+        messages = self._get_messages_from_request(request)
+
+        return {
+            "modelId": self.bedrock_model_id or model_id,
+            "inferenceConfig": {
+                "temperature": request.temperature,
+                "maxTokens": request.max_tokens,
+                "topP": request.top_p,
+            },
+            "messages": messages,
+        }
+
+    def make_request(self, request: Request) -> RequestResult:
+        raw_request = self.convert_request_to_raw_request(request)
+        cache_key = CachingClient.make_cache_key(raw_request, request)
+
+        def do_it() -> Dict[Any, Any]:
+            return self.bedrock_client.converse(**raw_request)
+
+        response, cached = self.cache.get(cache_key, do_it)
+
+        completions = self.convert_raw_response_to_completions(response, request)
+        dt = datetime.strptime(response["ResponseMetadata"]["HTTPHeaders"]["date"], "%a, %d %b %Y %H:%M:%S GMT")
+        # Use API reported latency rather than client measured latency
+        request_time = response["metrics"]["latencyMs"] / 1000
+
+        return RequestResult(
+            success=True,
+            cached=cached,
+            request_time=request_time,
+            request_datetime=int(dt.timestamp()),
+            completions=completions,
+            embedding=[],
+        )
+
+    def convert_raw_response_to_completions(self, response: Dict, request: Request) -> List[GeneratedOutput]:
+        completions: List[GeneratedOutput] = []
+        raw_completion = response["output"]
+        output_text = raw_completion["message"]["content"][0]["text"]
+        finish_reason = response["stopReason"]
+        completion = truncate_and_tokenize_response_text(
+            output_text.lstrip(), request, self.tokenizer, self.tokenizer_name, finish_reason
+        )
+        completions.append(completion)
+        return completions
+
+
+# Amazon Bedrock Client for Titan Models
+class BedrockTitanClient(BedrockClient):
+    _COMPLETION_REASON_TO_FINISH_REASON = {
+        "LENGTH": "length",
+        "FINISH": "endoftext",
+    }
+
+    # creator org for titan
+    model_provider = "amazon"
+
+    def convert_request_to_raw_request(self, request: Request) -> Dict:
+        # TODO: Support the following:
+        # - top_k_per_token
+        # - echo_prompt
+        # - num_completions
+        return {
+            "inputText": request.prompt,
+            "textGenerationConfig": {
+                "maxTokenCount": request.max_tokens,
+                # We ignore stop sequences in the request and always set stop sequences to the empty list.
+                # This is because:
+                #
+                # 1. The only permitted stop sequences are "|" and "User:"
+                #     - https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-titan-text.html
+                #     - https://github.com/boto/boto3/issues/3993
+                #     - https://github.com/aws/aws-sdk/issues/692
+                #
+                # 2. Titan has the tendency to emit "\n" as the first token in the generated text output,
+                #    which would cause the output to stop immediately if "\n" is in the stop_sequences.
+                "stopSequences": [],
+                "temperature": request.temperature,
+                "topP": request.top_p,
+            },
+        }
+
+    def convert_raw_response_to_completions(self, response: Dict, request: Request) -> List[GeneratedOutput]:
+        # TODO: Support the following:
+        # - tokens
+        # - logprob
+        completions: List[GeneratedOutput] = []
+
+        for raw_completion in response["results"]:
+            output_text = raw_completion["outputText"]
+            # Call lstrip() Titan has the tendency to emit "\n" as the first token in the generated text output.
+            finish_reason = BedrockTitanClient._COMPLETION_REASON_TO_FINISH_REASON.get(
+                raw_completion["completionReason"], raw_completion["completionReason"].lower()
+            )
+            completion = truncate_and_tokenize_response_text(
+                output_text.lstrip(), request, self.tokenizer, self.tokenizer_name, finish_reason
+            )
+            completions.append(completion)
+        return completions
+
+
+# Amazon Bedrock Client for Mistral Models
+class BedrockMistralClient(BedrockClient):
+    _COMPLETION_REASON_TO_FINISH_REASON = {
+        "length": "length",
+        "stop": "endoftext",
+    }
+
+    model_provider = "mistral"
+
+    def convert_request_to_raw_request(self, request: Request) -> Dict:
+        # TODO: Support the following:
+        # - top_k_per_token
+        # - echo_prompt
+        # - num_completions
+        return {
+            "prompt": f"[INST]{request.prompt}[/INST]",
+            "temperature": request.temperature,
+            "top_p": request.top_p,
+            "max_tokens": request.max_tokens,
+        }
+
+    def convert_raw_response_to_completions(self, response: Dict, request: Request) -> List[GeneratedOutput]:
+        # - logprob
+        completions: List[GeneratedOutput] = []
+
+        for raw_completion in response["outputs"]:
+            output_text = raw_completion["text"]
+
+            finish_reason = BedrockMistralClient._COMPLETION_REASON_TO_FINISH_REASON.get(
+                raw_completion["stop_reason"], raw_completion["stop_reason"].lower()
+            )
+            # Work around generated outputs with leading whitespace due to issue #2467
+            # TODO(#2467): Remove workaround
+            completion = truncate_and_tokenize_response_text(
+                output_text.lstrip(), request, self.tokenizer, self.tokenizer_name, finish_reason
+            )
+            completions.append(completion)
+
+        return completions
+
+
+# Amazon Bedrock Client for LLAMA Models
+class BedrockLlamaClient(BedrockClient):
+    _COMPLETION_REASON_TO_FINISH_REASON = {
+        "length": "length",
+        "stop": "endoftext",
+    }
+
+    model_provider = "meta"
+
+    def convert_request_to_raw_request(self, request: Request) -> Dict:
+        # TODO: Support the following:
+        # - top_k_per_token
+        # - echo_prompt
+        # - num_completions
+        return {
+            "prompt": f"[INST]{request.prompt}[/INST]",
+            "temperature": request.temperature,
+            "top_p": request.top_p,
+            "max_gen_len": request.max_tokens,
+        }
+
+    def convert_raw_response_to_completions(self, response: Dict, request: Request) -> List[GeneratedOutput]:
+        # - logprob
+        completions: List[GeneratedOutput] = []
+        output_text = response["generation"]
+
+        finish_reason = BedrockLlamaClient._COMPLETION_REASON_TO_FINISH_REASON.get(
+            response["stop_reason"], response["stop_reason"].lower()
+        )
+        # Work around generated outputs with leading whitespace due to issue #2467
+        # TODO(#2467): Remove workaround
+        completion = truncate_and_tokenize_response_text(
+            output_text.lstrip(), request, self.tokenizer, self.tokenizer_name, finish_reason
+        )
+        completions.append(completion)
+
+        return completions
+
+
+class BedrockPalmyraClient(BedrockClient):
+    """Amazon Bedrock model for Palmyra models"""
+
+    _COMPLETION_REASON_TO_FINISH_REASON = {
+        "length": "length",
+        "stop": "endoftext",
+    }
+
+    model_provider = "writer"
+
+    def _get_messages_from_request(self, request: Request) -> List[Dict]:
+        if request.prompt and request.messages:
+            raise NonRetriableException(f"Only one of `prompt` and `messages` may be set in request: {request}")
+        if request.multimodal_prompt:
+            raise NonRetriableException("`multimodal_prompt` is not supported by WriterClient")
+        if request.messages:
+            return [{"role": message["role"], "content": message["content"]} for message in request.messages]
+        else:
+            return [{"role": "user", "content": request.prompt}]
+
+    def convert_request_to_raw_request(self, request: Request) -> Dict:
+        return {
+            "messages": self._get_messages_from_request(request),
+            "temperature": request.temperature,
+            "top_p": request.top_p,
+            "max_tokens": request.max_tokens,
+            "stop": request.stop_sequences,
+            "seed": request.random,
+            "presence_penalty": request.presence_penalty,
+            "frequency_penalty": request.frequency_penalty,
+        }
+
+    def convert_raw_response_to_completions(self, response: Dict, request: Request) -> List[GeneratedOutput]:
+        completions: List[GeneratedOutput] = []
+        if response.get("object") == "error":
+            raise Exception(f"BedrockPalmyraClient request failed: {response['message']}")
+        for choice in response["choices"]:
+            assert choice["message"]["role"] == "assistant"
+            output_text = choice["message"]["content"]
+            finish_reason = BedrockPalmyraClient._COMPLETION_REASON_TO_FINISH_REASON.get(
+                choice["finish_reason"], choice["finish_reason"].lower()
+            )
+            completion = truncate_and_tokenize_response_text(
+                output_text, request, self.tokenizer, self.tokenizer_name, finish_reason
+            )
+            completions.append(completion)
+        return completions

helm/clients/bedrock_utils.py

@@ -0,0 +1,105 @@
+"""Helper utilities for working with Amazon Bedrock."""
+
+import os
+from typing import Optional
+
+from helm.common.hierarchical_logger import hlog
+from helm.common.optional_dependencies import handle_module_not_found_error
+
+try:
+    import boto3
+    from boto3 import Session
+    from botocore.config import Config
+except ModuleNotFoundError as e:
+    handle_module_not_found_error(e, ["aws"])
+
+
+# From https://github.com/aws-samples/amazon-bedrock-workshop/blob/main/01_Generation/00_generate_w_bedrock.ipynb
+# MIT-0 Licensed
+def get_bedrock_client(
+    assumed_role: Optional[str] = None,
+    region: Optional[str] = None,
+    runtime: Optional[bool] = True,
+):
+    """Create a boto3 client for Amazon Bedrock, with optional configuration overrides
+
+    Parameters
+    ----------
+    assumed_role :
+        Optional ARN of an AWS IAM role to assume for calling the Bedrock service. If not
+        specified, the current active credentials will be used.
+    region :
+        Optional name of the AWS Region in which the service should be called (e.g. "us-east-1").
+        If not specified, AWS_REGION or AWS_DEFAULT_REGION environment variable will be used.
+    runtime :
+        Optional choice of getting different client to perform operations with the Amazon Bedrock service.
+    """
+    if region is None:
+        target_region = os.environ.get("AWS_REGION", os.environ.get("AWS_DEFAULT_REGION"))
+    else:
+        target_region = region
+
+    session_kwargs = {"region_name": target_region}
+    client_kwargs = {**session_kwargs}
+
+    profile_name = os.environ.get("AWS_PROFILE")
+    if profile_name:
+        session_kwargs["profile_name"] = profile_name
+
+    retry_config = Config(
+        region_name=target_region,
+        retries={
+            "max_attempts": 10,
+            "mode": "standard",
+        },
+    )
+    session = boto3.Session(**session_kwargs)
+
+    if assumed_role:
+        sts = session.client("sts")
+        response = sts.assume_role(RoleArn=str(assumed_role), RoleSessionName="crfm-helm")
+        client_kwargs["aws_access_key_id"] = response["Credentials"]["AccessKeyId"]
+        client_kwargs["aws_secret_access_key"] = response["Credentials"]["SecretAccessKey"]
+        client_kwargs["aws_session_token"] = response["Credentials"]["SessionToken"]
+
+    if runtime:
+        service_name = "bedrock-runtime"
+    else:
+        service_name = "bedrock"
+
+    bedrock_client = session.client(service_name=service_name, config=retry_config, **client_kwargs)
+
+    hlog(f"Amazon Bedrock client successfully created with endpoint {bedrock_client._endpoint}")
+    return bedrock_client
+
+
+def get_bedrock_client_v1(
+    region: Optional[str] = None,
+    service_name: str = "bedrock-runtime",
+    assumed_role: Optional[str] = None,
+    read_timeout: int = 5000,
+    connect_timeout: int = 5000,
+    max_attempts: int = 10,
+):
+    boto_config = Config(
+        read_timeout=read_timeout, connect_timeout=connect_timeout, retries={"max_attempts": max_attempts}
+    )
+
+    if assumed_role:
+        session = boto3.Session(region_name=region)
+        # Assume role and get credentials
+        sts = session.client("sts")
+        creds = sts.assume_role(RoleArn=str(assumed_role), RoleSessionName="crfm-helm")["Credentials"]
+        session = Session(
+            aws_access_key_id=creds["AccessKeyId"],
+            aws_secret_access_key=creds["SecretAccessKey"],
+            aws_session_token=creds["SessionToken"],
+        )
+        return session.client(
+            service_name=service_name,
+            region_name=region,
+            config=boto_config,
+        )
+
+    # default to instance role to get the aws credentials or aws configured credentials
+    return boto3.client(service_name=service_name, region_name=region, config=boto_config)

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

@@ -1,11 +1,13 @@
 import json
 from abc import ABC, abstractmethod
-from typing import Dict, List, Optional
+from typing import List, Mapping, Optional, cast
 
-from helm.common.hierarchical_logger import hlog
+from helm.common.hierarchical_logger import hwarn
 from helm.common.media_object import MultimediaObject, TEXT_TYPE
-from helm.common.request import Request, RequestResult, Sequence, Token
+from helm.common.request import Request, RequestResult, GeneratedOutput, Token
 from helm.common.cache import Cache, CacheConfig
+from helm.common.tokenization_request import DecodeRequest, TokenizationRequest
+from helm.tokenizers.tokenizer import Tokenizer
 
 
 class Client(ABC):
@@ -30,24 +32,32 @@ class CachingClient(Client):
         self.cache = Cache(cache_config) if cache_config is not None else None
 
     @staticmethod
-    def make_cache_key(raw_request: Dict, request: Request) -> Dict:
+    def make_cache_key(raw_request: Mapping, request: Request) -> Mapping:
         """
         Construct the key for the cache using the raw request.
         Add `request.random` to the key, if defined.
         """
         if request.random is not None:
             assert "random" not in raw_request
-            cache_key = {**raw_request, "random": request.random}
+            return {**raw_request, "random": request.random}
         else:
-            cache_key = raw_request
-        return cache_key
+            return {**raw_request}
 
 
-def truncate_sequence(sequence: Sequence, request: Request, print_warning: bool = True) -> Sequence:
+def truncate_sequence(
+    sequence: GeneratedOutput,
+    request: Request,
+    end_of_text_token: Optional[str] = None,
+    print_warning: bool = True,
+) -> GeneratedOutput:
     """
     Certain providers have bugs where they aren't respecting max_tokens,
     stop_sequences and the end of text token, so as a hack, we have to manually
     truncate the suffix of `sequence` and `tokens` as a post-hoc process.
+
+    This method is unsafe and may produce warnings or incorrect results.
+    Prefer using the safer truncate_and_tokenize_response_text() method instead
+    if your use case satisfies its requirements.
     """
     # TODO: if echo_prompt, then we should only ignore the prompt, but we don't
     # know how many tokens the prompt takes up.
@@ -55,10 +65,14 @@ def truncate_sequence(sequence: Sequence, request: Request, print_warning: bool
     # where max_tokens = 0, so there's nothing to truncate.
     if request.echo_prompt:
         if request.max_tokens != 0:
-            hlog("WARNING: don't know how to handle echo_prompt and max_tokens > 0, not truncating")
+            hwarn("don't know how to handle echo_prompt and max_tokens > 0, not truncating")
         return sequence
 
-    for stop in request.stop_sequences:
+    if end_of_text_token:
+        stop_sequences = request.stop_sequences + [end_of_text_token]
+    else:
+        stop_sequences = request.stop_sequences
+    for stop in stop_sequences:
         # Find `stop` in the text
         try:
             new_text = sequence.text[: sequence.text.index(stop)]
@@ -76,8 +90,8 @@ def truncate_sequence(sequence: Sequence, request: Request, print_warning: bool
             new_tokens.append(token)
 
         if len(new_text) < len(sequence.text) and len(new_tokens) == len(sequence.tokens):
-            hlog(
-                f"WARNING: Stripped characters from text ({len(sequence.text)} -> {len(new_text)}), "
+            hwarn(
+                f"Stripped characters from text ({len(sequence.text)} -> {len(new_text)}), "
                 f"but wasn't able to strip the tokens"
             )
 
@@ -85,14 +99,14 @@ def truncate_sequence(sequence: Sequence, request: Request, print_warning: bool
         new_logprob = sum(token.logprob for token in new_tokens)
 
         if print_warning:
-            hlog(f"WARNING: truncate_sequence needs to strip {json.dumps(stop)}")
+            hwarn(f"truncate_sequence needs to strip {json.dumps(stop)}")
 
-        sequence = Sequence(text=new_text, logprob=new_logprob, tokens=new_tokens)
+        sequence = GeneratedOutput(text=new_text, logprob=new_logprob, tokens=new_tokens)
 
     # Truncate based on the max number of tokens.
     if len(sequence.tokens) > request.max_tokens:
         if print_warning:
-            hlog(f"WARNING: truncate_sequence needs to truncate {len(sequence.tokens)} down to {request.max_tokens}")
+            hwarn(f"truncate_sequence needs to truncate {len(sequence.tokens)} down to {request.max_tokens}")
         new_tokens = sequence.tokens[: request.max_tokens]
 
         # This is imperfect stitching together of tokens, so just to make sure this is okay
@@ -100,15 +114,76 @@ def truncate_sequence(sequence: Sequence, request: Request, print_warning: bool
         # Usually, in our benchmark, max_tokens is active when it's 1, so hopefully this isn't an issue.
         new_text = "".join(token.text for token in new_tokens)
         if not sequence.text.startswith(new_text):
-            hlog(f"WARNING: {json.dumps(sequence.text)} does not start with truncated text {json.dumps(new_text)}")
+            hwarn(f"{json.dumps(sequence.text)} does not start with truncated text {json.dumps(new_text)}")
 
         new_logprob = sum(token.logprob for token in new_tokens)
 
-        sequence = Sequence(text=new_text, logprob=new_logprob, tokens=new_tokens)
+        sequence = GeneratedOutput(text=new_text, logprob=new_logprob, tokens=new_tokens)
 
     return sequence
 
 
+def truncate_and_tokenize_response_text(
+    text: str,
+    request: Request,
+    tokenizer: Tokenizer,
+    tokenizer_name: str,
+    end_of_text_token: Optional[str] = None,
+    original_finish_reason: str = "endoftext",
+) -> GeneratedOutput:
+    """Truncate a string-only response to respect stop_sequences and max_tokens.
+
+    This can only be used if all of the following conditions are true:
+
+    - You have access to the tokenizer.
+    - The request has echo_prompt = False.
+    - The tokenizer supports encoding and decoding.
+    - The tokenizer's tokenize() method supports truncation.
+    - The model's response is text-only.
+    - The model's response not already provide the tokenized text.
+    - The model's response does not provide logprobs.
+
+    This method is safer than truncate_sequence() and should be preferred if the above conditions are met.
+    Unlike truncate_sequence(), this method will not produce warnings or incorrect results.
+    This is because the the tokens are derived from the truncated text using the tokenizer,
+    so the text and the tokens in the resulting result are guranteed to match."""
+    # Finish reason strings are token from basic_metrics._compute_finish_reason_metrics()
+    finish_reason: str = original_finish_reason
+    if request.echo_prompt:
+        raise Exception("truncate_and_tokenize_response_text() does not support requests with echo_prompt = True")
+
+    if end_of_text_token:
+        stop_sequences = request.stop_sequences + [end_of_text_token]
+    else:
+        stop_sequences = request.stop_sequences
+    for stop_sequence in stop_sequences:
+        try:
+            text = text[: text.index(stop_sequence)]
+            finish_reason = "stop"
+        except ValueError:
+            pass
+
+    token_strings = cast(
+        List[str], tokenizer.tokenize(TokenizationRequest(text=text, tokenizer=tokenizer_name)).raw_tokens
+    )
+    if len(token_strings) > request.max_tokens:
+        encoded_ints = cast(
+            List[int],
+            tokenizer.tokenize(
+                TokenizationRequest(
+                    text=text, tokenizer=tokenizer_name, encode=True, truncation=True, max_length=request.max_tokens
+                )
+            ).raw_tokens,
+        )
+        text = tokenizer.decode(DecodeRequest(encoded_ints, tokenizer_name)).text
+        token_strings = cast(
+            List[str], tokenizer.tokenize(TokenizationRequest(text=text, tokenizer=tokenizer_name)).raw_tokens
+        )
+        finish_reason = "length"
+    tokens = [Token(text=token_string, logprob=0.0) for token_string in token_strings]
+    return GeneratedOutput(text=text, logprob=0.0, tokens=tokens, finish_reason={"reason": finish_reason})
+
+
 def cleanup_str(token: str, tokenizer_name: Optional[str] = None) -> str:
     """
     Certain tokenizers introduce special characters to represent spaces, such as