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

helm/clients/vertexai_client.py

@@ -0,0 +1,488 @@
+import requests
+from abc import ABC, abstractmethod
+from threading import Lock
+from typing import Any, Dict, Mapping, Optional, List, Union, cast
+
+from helm.common.cache import CacheConfig
+from helm.common.hierarchical_logger import hexception
+from helm.common.multimodal_request_utils import get_contents_as_bytes
+from helm.common.media_object import TEXT_TYPE
+from helm.common.optional_dependencies import handle_module_not_found_error
+from helm.common.request import wrap_request_time, Request, RequestResult, GeneratedOutput, ErrorFlags
+from helm.clients.client import CachingClient, truncate_sequence, generate_uid_for_multimodal_prompt
+
+try:
+    import vertexai
+    from vertexai.language_models import TextGenerationModel, TextGenerationResponse  # PaLM2
+    from vertexai.preview.generative_models import (
+        GenerativeModel,
+        GenerationResponse,
+        Candidate,
+        Content,
+        Part,
+        Image,
+    )  # Gemini
+    from google.cloud.aiplatform_v1beta1.types import SafetySetting, HarmCategory
+except ModuleNotFoundError as e:
+    handle_module_not_found_error(e, ["google"])
+
+
+_models_lock: Lock = Lock()
+_models: Dict[str, Any] = {}
+
+
+class VertexAIContentBlockedError(Exception):
+    pass
+
+
+class SafetySettingPresets:
+    BLOCK_NONE = "block_none"  # Disable all blocking
+    DEFAULT = "default"  # Use default safety settings
+
+
+def _get_safety_settings_for_preset(
+    safety_settings_preset: Optional[str],
+) -> Optional[Dict[HarmCategory, SafetySetting.HarmBlockThreshold]]:
+    """Get the safety settings for the safety_settings_preset.
+
+    If safety_settings_preset is None, use the default value of BLOCK_NONE (*not* DEFAULT)."""
+    if safety_settings_preset is None or safety_settings_preset == SafetySettingPresets.BLOCK_NONE:
+        return {
+            harm_category: SafetySetting.HarmBlockThreshold(SafetySetting.HarmBlockThreshold.BLOCK_NONE)
+            for harm_category in iter(HarmCategory)
+        }
+    elif safety_settings_preset == SafetySettingPresets.DEFAULT:
+        return None
+    else:
+        raise ValueError(f"Unknown safety_settings_preset: {safety_settings_preset}")
+
+
+class VertexAIClient(CachingClient, ABC):
+    """Client for Vertex AI models"""
+
+    def __init__(
+        self,
+        cache_config: CacheConfig,
+        project_id: str,
+        location: str,
+        safety_settings_preset: Optional[str] = None,
+        vertexai_model: Optional[str] = None,
+    ) -> None:
+        super().__init__(cache_config=cache_config)
+        self.project_id = project_id
+        self.location = location
+
+        self.safety_settings_preset = safety_settings_preset
+        self.safety_settings = _get_safety_settings_for_preset(safety_settings_preset)
+
+        self.vertexai_model = vertexai_model
+
+        vertexai.init(project=self.project_id, location=self.location)
+
+    def _get_model_name_for_request(self, request: Request) -> str:
+        if self.vertexai_model is not None:
+            return self.vertexai_model
+        return request.model_engine
+
+    def make_cache_key_with_safety_settings_preset(self, raw_request: Mapping, request: Request) -> Mapping:
+        """Construct the key for the cache using the raw request.
+
+        Add `self.safety_settings_preset` to the key, if not None."""
+        if self.safety_settings_preset is not None:
+            assert "safety_settings_preset" not in raw_request
+            return {
+                **CachingClient.make_cache_key(raw_request, request),
+                "safety_settings_preset": self.safety_settings_preset,
+            }
+        else:
+            return CachingClient.make_cache_key(raw_request, request)
+
+    @abstractmethod
+    def make_request(self, request: Request) -> RequestResult:
+        raise NotImplementedError
+
+
+class VertexAITextClient(VertexAIClient):
+    """Client for Vertex AI text models
+    This client is used for PaLM2 for example."""
+
+    def make_request(self, request: Request) -> RequestResult:
+        """Make a request"""
+        parameters: Dict[str, Any] = {
+            "temperature": request.temperature,
+            "max_output_tokens": request.max_tokens,
+            "top_k": request.top_k_per_token,
+            "top_p": request.top_p,
+            "stop_sequences": request.stop_sequences,
+            "candidate_count": request.num_completions,
+            # TODO #2084: Add support for these parameters.
+            # The parameters "echo", "frequency_penalty", and "presence_penalty" are supposed to be supported
+            # in an HTTP request (See https://cloud.google.com/vertex-ai/docs/generative-ai/model-reference/text),
+            # but they are not supported in the Python SDK:
+            # https://github.com/googleapis/python-aiplatform/blob/beae48f63e40ea171c3f1625164569e7311b8e5a/vertexai/language_models/_language_models.py#L968C1-L980C1
+            # "frequency_penalty": request.frequency_penalty,
+            # "presence_penalty": request.presence_penalty,
+            # "echo": request.echo_prompt,
+        }
+
+        completions: List[GeneratedOutput] = []
+        model_name: str = self._get_model_name_for_request(request)
+
+        try:
+
+            def do_it() -> Dict[str, Any]:
+                model = TextGenerationModel.from_pretrained(model_name)
+                response = model.predict(request.prompt, **parameters)
+                candidates: List[TextGenerationResponse] = response.candidates
+                response_dict = {
+                    "predictions": [{"text": completion.text for completion in candidates}],
+                }  # TODO: Extract more information from the response
+                return response_dict
+
+            # 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.
+            # Same for the prompt.
+            cache_key = self.make_cache_key_with_safety_settings_preset(
+                {
+                    "engine": model_name,
+                    "prompt": request.prompt,
+                    **parameters,
+                },
+                request,
+            )
+
+            response, cached = self.cache.get(cache_key, wrap_request_time(do_it))
+        except (requests.exceptions.RequestException, AssertionError) as e:
+            hexception(e)
+            error: str = f"VertexAITextClient error: {e}"
+            return RequestResult(success=False, cached=False, error=error, completions=[], embedding=[])
+
+        for prediction in response["predictions"]:
+            response_text = prediction["text"]
+
+            # The Python SDK does not support echo
+            text: str = request.prompt + response_text if request.echo_prompt else response_text
+
+            # TODO #2085: Add support for log probs.
+            # Once again, log probs seem to be supported by the API but not by the Python SDK.
+            # HTTP Response body reference:
+            # https://cloud.google.com/vertex-ai/docs/generative-ai/model-reference/text#response_body
+            # Python SDK reference:
+            # https://github.com/googleapis/python-aiplatform/blob/beae48f63e40ea171c3f1625164569e7311b8e5a/vertexai/language_models/_language_models.py#L868
+            completion = GeneratedOutput(text=text, logprob=0, tokens=[])
+            sequence = truncate_sequence(completion, request, print_warning=True)
+            completions.append(sequence)
+
+        return RequestResult(
+            success=True,
+            cached=cached,
+            request_time=response["request_time"],
+            request_datetime=response["request_datetime"],
+            completions=completions,
+            embedding=[],
+        )
+
+
+class VertexAIChatClient(VertexAIClient):
+    """Client for Vertex AI chat models (e.g., Gemini). Supports multimodal prompts."""
+
+    # Enum taken from:
+    # https://cloud.google.com/vertex-ai/docs/reference/rpc/google.cloud.aiplatform.v1beta1#google.cloud.aiplatform.v1beta1.Candidate.FinishReason
+    # We don't directly import this enum because it can differ between different Vertex AI library versions.
+    CONTENT_BLOCKED_FINISH_REASONS: List[int] = [
+        3,  # SAFETY
+        4,  # RECITATION
+        6,  # BLOCKLIST
+        7,  # PROHIBITED_CONTENT
+        8,  # SPII (Sensitive Personally Identifiable Information)
+    ]
+
+    @staticmethod
+    def get_model(model_name: str) -> GenerativeModel:
+        global _models_lock
+        global _models
+
+        with _models_lock:
+            if model_name not in _models:
+                _models[model_name] = GenerativeModel(model_name)
+            return _models[model_name]
+
+    def make_request(self, request: Request) -> RequestResult:
+        """Make a request"""
+        # mypy is unhappy without this cast
+        contents: Union[List[Union[str, Image, Part]], List[Content]] = cast(
+            List[Union[str, Image, Part]], [request.prompt]
+        )
+
+        # For the multimodal case, build up the content with the media objects of `request.multimodal_prompt`
+        if request.multimodal_prompt is not None:
+            return self._make_multimodal_request(request)
+
+        if request.messages is not None:
+            role_mapping = {"user": "user", "assistant": "model"}
+            contents = [
+                Content(role=role_mapping.get(msg["role"], "user"), parts=[Part.from_text(msg["content"])])
+                for msg in request.messages
+            ]
+
+        parameters: Dict[str, Any] = {
+            "temperature": request.temperature,
+            "max_output_tokens": request.max_tokens,
+            "top_k": request.top_k_per_token,
+            "top_p": request.top_p,
+            "stop_sequences": request.stop_sequences,
+            "candidate_count": request.num_completions,
+            # TODO #2084: Add support for these parameters.
+            # The parameters "echo", "frequency_penalty", and "presence_penalty" are supposed to be supported
+            # in an HTTP request (See https://cloud.google.com/vertex-ai/docs/generative-ai/model-reference/text),
+            # but they are not supported in the Python SDK:
+            # https://github.com/googleapis/python-aiplatform/blob/beae48f63e40ea171c3f1625164569e7311b8e5a/vertexai/language_models/_language_models.py#L968C1-L980C1
+            # "frequency_penalty": request.frequency_penalty,
+            # "presence_penalty": request.presence_penalty,
+            # "echo": request.echo_prompt,
+        }
+
+        completions: List[GeneratedOutput] = []
+        model_name: str = self._get_model_name_for_request(request)
+        model = self.get_model(model_name)
+
+        try:
+
+            def do_it() -> Dict[str, Any]:
+                # Here we differ from Vertex AI's tutorial.
+                # https://cloud.google.com/vertex-ai/docs/generative-ai/multimodal/send-chat-prompts-gemini#send_chat_prompts   # noqa: E501
+                # It would advise to use model.start_chat() but since we do not want to use Chat capabilities of
+                # Vertex AI, we use model.generate_text() instead. Furthermore, chat.send_message() restricts the
+                # output to only one candidate.
+                # chat: ChatSession = model.start_chat()
+                # See: https://github.com/googleapis/python-aiplatform/blob/e8c505751b10a9dc91ae2e0d6d13742d2abf945c/vertexai/generative_models/_generative_models.py#L812  # noqa: E501
+                response: GenerationResponse = model.generate_content(
+                    contents, generation_config=parameters, safety_settings=self.safety_settings
+                )
+                candidates: List[Candidate] = response.candidates
+
+                # Depending on the version of the Vertex AI library and the type of prompt blocking,
+                # prompt blocking can show up in many ways, so this defensively handles most of these ways
+                if response.prompt_feedback and response.prompt_feedback.block_reason:
+                    raise VertexAIContentBlockedError(
+                        f"Prompt blocked with reason: {response.prompt_feedback.block_reason}"
+                    )
+                if not candidates:
+                    raise VertexAIContentBlockedError(f"No candidates in response: {response}")
+                predictions: List[Dict[str, Any]] = []
+                for candidate in candidates:
+                    # Depending on the version of the Vertex AI library and the type of prompt blocking,
+                    # content blocking can show up in many ways, so this defensively handles most of these ways
+                    if candidate.finish_reason in VertexAIChatClient.CONTENT_BLOCKED_FINISH_REASONS:
+                        raise VertexAIContentBlockedError(f"Content blocked with reason: {candidate.finish_reason}")
+                    if not candidate.content:
+                        raise VertexAIContentBlockedError(f"No content in candidate: {candidate}")
+                    if not candidate.content.parts:
+                        if candidate.finish_reason == 2:  # MAX_TOKENS
+                            # This means that there is no text output because the maximum number of tokens were
+                            # reached during thinking.
+                            predictions.append({"text": ""})
+                        else:
+                            raise VertexAIContentBlockedError(f"No content parts in candidate: {candidate}")
+                    else:
+                        predictions.append({"text": candidate.content.text})
+                    # TODO: Extract more information from the response
+                return {"predictions": predictions}
+
+            # 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.
+            # Same for the prompt.
+            cache_key = self.make_cache_key_with_safety_settings_preset(
+                {
+                    "model_name": model_name,
+                    "prompt": request.messages or request.prompt,
+                    **parameters,
+                },
+                request,
+            )
+
+            response, cached = self.cache.get(cache_key, wrap_request_time(do_it))
+        except VertexAIContentBlockedError as e:
+            return RequestResult(
+                success=False,
+                cached=False,
+                error=f"Content blocked: {str(e)}",
+                completions=[],
+                embedding=[],
+                error_flags=ErrorFlags(is_retriable=False, is_fatal=False),
+            )
+        except (requests.exceptions.RequestException, AssertionError) as e:
+            hexception(e)
+            error: str = f"VertexAITextClient error: {e}"
+            return RequestResult(success=False, cached=False, error=error, completions=[], embedding=[])
+
+        # Handle cached responses with blocked content from old versions of HELM.
+        if response["predictions"] is None:
+            return RequestResult(
+                success=False,
+                cached=False,
+                error=f"Content blocked error in cached response: {str(response)}",
+                completions=[],
+                embedding=[],
+                error_flags=ErrorFlags(is_retriable=False, is_fatal=False),
+                request_time=response["request_time"],
+                request_datetime=response["request_datetime"],
+            )
+
+        for prediction in response["predictions"]:
+            # Handle cached responses with blocked content from old versions of HELM.
+            if "text" not in prediction:
+                return RequestResult(
+                    success=False,
+                    cached=False,
+                    error=f"Content blocked error in cached prediction: {str(prediction)}",
+                    completions=[],
+                    embedding=[],
+                    error_flags=ErrorFlags(is_retriable=False, is_fatal=False),
+                    request_time=response["request_time"],
+                    request_datetime=response["request_datetime"],
+                )
+            response_text = prediction["text"]
+
+            # The Python SDK does not support echo
+            text: str = request.prompt + response_text if request.echo_prompt else response_text
+            completion = GeneratedOutput(text=text, logprob=0, tokens=[])
+            sequence = truncate_sequence(completion, request, print_warning=True)
+            completions.append(sequence)
+
+        return RequestResult(
+            success=True,
+            cached=cached,
+            request_time=response["request_time"],
+            request_datetime=response["request_datetime"],
+            completions=completions,
+            embedding=[],
+        )
+
+    def _make_multimodal_request(self, request: Request) -> RequestResult:
+        # Contents can either be text or a list of multimodal content made up of text, images or other content
+        contents: Union[str, List[Union[str, Any]]] = request.prompt
+        # Used to generate a unique cache key for this specific request
+        assert request.multimodal_prompt is not None
+        prompt_key: str = generate_uid_for_multimodal_prompt(request.multimodal_prompt)
+
+        # For the multimodal case, build up the content with the media objects of `request.multimodal_prompt`
+        contents = []
+        for media_object in request.multimodal_prompt.media_objects:
+            if media_object.is_type("image") and media_object.location:
+                contents.append(Part.from_image(Image.load_from_file(media_object.location)))
+            elif media_object.is_type("video") and media_object.location:
+                # Following this example
+                # https://cloud.google.com/vertex-ai/generative-ai/docs/samples/googlegenaisdk-textgen-with-local-video
+                with open(media_object.location, "rb") as fp:
+                    video_content = fp.read()
+                contents.append(Part.from_data(data=video_content, mime_type=media_object.content_type))
+            elif media_object.is_type("audio") and media_object.location:
+                contents.append(
+                    Part.from_data(get_contents_as_bytes(media_object.location), mime_type=media_object.content_type)
+                )
+            elif media_object.is_type(TEXT_TYPE):
+                if media_object.text is None:
+                    raise ValueError("MediaObject of text type has missing text field value")
+                contents.append(media_object.text)
+            else:
+                raise ValueError(f"Unrecognized MediaObject type {media_object.type}")
+
+        parameters = {
+            "temperature": request.temperature,
+            "max_output_tokens": request.max_tokens,
+            "top_k": request.top_k_per_token,
+            "top_p": request.top_p,
+            "stop_sequences": request.stop_sequences,
+            "candidate_count": 1,
+        }
+
+        completions: List[GeneratedOutput] = []
+        model_name: str = self._get_model_name_for_request(request)
+        model = self.get_model(model_name)
+
+        request_time = 0
+        request_datetime: Optional[int] = None
+        all_cached = True
+
+        # Gemini Vision only supports generating 1-2 candidates at a time, so make `request.num_completions` requests
+        for completion_index in range(request.num_completions):
+            try:
+
+                def do_it() -> Dict[str, Any]:
+                    response: GenerationResponse = model.generate_content(
+                        contents, generation_config=parameters, safety_settings=self.safety_settings
+                    )
+                    # Depending on the version of the Vertex AI library and the type of prompt blocking,
+                    # prompt blocking can show up in many ways, so this defensively handles most of these ways
+                    if response.prompt_feedback and response.prompt_feedback.block_reason:
+                        raise VertexAIContentBlockedError(
+                            f"Prompt blocked with reason: {response.prompt_feedback.block_reason}"
+                        )
+                    if not response.candidates:
+                        raise VertexAIContentBlockedError(f"No candidates in response: {response}")
+                    # We should only have one candidate
+                    assert (
+                        len(response.candidates) == 1
+                    ), f"Expected 1 candidate since candidate_count is 1, got {len(response.candidates)}."
+                    candidate = response.candidates[0]
+                    # Depending on the version of the Vertex AI library and the type of prompt blocking,
+                    # content blocking can show up in many ways, so this defensively handles most of these ways
+                    if candidate.finish_reason in VertexAIChatClient.CONTENT_BLOCKED_FINISH_REASONS:
+                        raise VertexAIContentBlockedError(f"Content blocked with reason: {candidate.finish_reason}")
+                    if not candidate.content:
+                        raise VertexAIContentBlockedError(f"No content in candidate: {candidate}")
+                    if not candidate.content.parts:
+                        raise VertexAIContentBlockedError(f"No content parts in candidate: {candidate}")
+                    return {"predictions": [{"text": candidate.text}]}
+
+                raw_cache_key = {"model_name": model_name, "prompt": prompt_key, **parameters}
+                if completion_index > 0:
+                    raw_cache_key["completion_index"] = completion_index
+
+                cache_key = self.make_cache_key_with_safety_settings_preset(raw_cache_key, request)
+                response, cached = self.cache.get(cache_key, wrap_request_time(do_it))
+            except requests.exceptions.RequestException as e:
+                hexception(e)
+                error: str = f"Gemini Vision error: {e}"
+                return RequestResult(success=False, cached=False, error=error, completions=[], embedding=[])
+            except VertexAIContentBlockedError as e:
+                return RequestResult(
+                    success=False,
+                    cached=False,
+                    error=f"Content blocked: {str(e)}",
+                    completions=[],
+                    embedding=[],
+                    error_flags=ErrorFlags(is_retriable=False, is_fatal=False),
+                )
+
+            if "error" in response:
+                return RequestResult(
+                    success=False,
+                    cached=True,
+                    error=f"Content blocked error in cached response: {str(response)}",
+                    completions=[],
+                    embedding=[],
+                    error_flags=ErrorFlags(is_retriable=False, is_fatal=False),
+                    request_time=response["request_time"],
+                    request_datetime=response["request_datetime"],
+                )
+
+            response_text = response["predictions"][0]["text"]
+            completion = GeneratedOutput(text=response_text, logprob=0, tokens=[])
+            sequence = truncate_sequence(completion, request, print_warning=True)
+            completions.append(sequence)
+
+            request_time += response["request_time"]
+            # Use the datetime from the first completion because that's when the request was fired
+            request_datetime = request_datetime or response.get("request_datetime")
+            all_cached = all_cached and cached
+
+        return RequestResult(
+            success=True,
+            cached=all_cached,
+            request_time=request_time,
+            request_datetime=request_datetime,
+            completions=completions,
+            embedding=[],
+        )

helm/clients/vision_language/huggingface_vision2seq_client.py

@@ -0,0 +1,148 @@
+from threading import Lock
+from typing import Any, Dict, List, Optional
+
+from dataclasses import dataclass
+from transformers import AutoProcessor, AutoModelForVision2Seq
+from transformers.image_utils import load_image
+import torch
+
+from helm.common.cache import CacheConfig
+from helm.common.gpu_utils import get_torch_device_name, is_cuda_available
+from helm.common.hierarchical_logger import hexception, hlog, htrack_block
+from helm.common.media_object import TEXT_TYPE
+from helm.common.request import Request, RequestResult, GeneratedOutput, Token
+from helm.common.request import wrap_request_time
+from helm.common.tokenization_request import TokenizationRequest
+from helm.clients.client import CachingClient, generate_uid_for_multimodal_prompt
+from helm.tokenizers.tokenizer import Tokenizer
+
+
+@dataclass(frozen=True)
+class Vision2SeqModelProcessor:
+    """Loaded model and processor."""
+
+    model: AutoModelForVision2Seq
+    processor: AutoProcessor
+
+
+_models_lock: Lock = Lock()
+_models: Dict[str, Optional[Vision2SeqModelProcessor]] = {
+    "HuggingFaceM4/idefics2-8b": None,
+}
+
+
+class HuggingFaceVision2SeqClient(CachingClient):
+    """
+    Models for Vision2Seq models from HuggingFace.
+    """
+
+    ASSISTANT_PREFIX: str = "Assistant:"
+
+    def __init__(self, tokenizer: Tokenizer, tokenizer_name: str, cache_config: CacheConfig):
+        super().__init__(cache_config=cache_config)
+        self.tokenizer = tokenizer
+        self.tokenizer_name = tokenizer_name
+        self._device: str = get_torch_device_name()
+
+    def _get_model(self, checkpoint: str) -> Vision2SeqModelProcessor:
+        global _models_lock
+        global _models
+
+        # Ensure that only one thread is loading the model at a time
+        with _models_lock:
+            loaded_model_processor = _models[checkpoint]
+            if loaded_model_processor is None:
+                hlog(f"Loading model {checkpoint} and caching in memory...")
+                torch_dtype: torch.dtype = torch.float16 if is_cuda_available() else torch.float32
+                model = AutoModelForVision2Seq.from_pretrained(checkpoint, torch_dtype=torch_dtype).to(self._device)
+                processor = AutoProcessor.from_pretrained(checkpoint)
+
+                _models[checkpoint] = Vision2SeqModelProcessor(model, processor)
+                loaded_model_processor = _models[checkpoint]
+
+        assert loaded_model_processor is not None
+        return loaded_model_processor
+
+    def make_request(self, request: Request) -> RequestResult:
+        assert request.model_deployment in _models, f"Not a valid model for this client: {request.model_deployment}"
+        assert request.multimodal_prompt is not None, "Multimodal prompt is required"
+
+        loaded_model_processor: Vision2SeqModelProcessor = self._get_model(request.model_deployment)
+        model = loaded_model_processor.model
+        processor = loaded_model_processor.processor
+
+        generation_args: Dict[str, Any] = {
+            "max_new_tokens": request.max_tokens,
+        }
+
+        image_paths: List[str] = []
+        multimodal_prompt: List[Dict[str, str]] = []
+        for media_object in request.multimodal_prompt.media_objects:
+            if media_object.is_type("image") and media_object.location:
+                image_paths.append(media_object.location)
+                multimodal_prompt.append({"type": "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")
+
+                multimodal_prompt.append({"type": "text", "text": media_object.text})
+            else:
+                raise ValueError(f"Unrecognized MediaObject type {media_object.type}")
+
+        completions: List[GeneratedOutput] = []
+        with htrack_block(f"Generating for prompt: {request.multimodal_prompt.text}"):
+            try:
+
+                def do_it() -> Dict[str, Any]:
+                    messages = [{"role": "user", "content": multimodal_prompt}]
+                    prompt = processor.apply_chat_template(messages, add_generation_prompt=True)  # type: ignore
+                    inputs = processor(  # type: ignore
+                        text=[prompt] * request.num_completions,
+                        images=[
+                            [load_image(image_path) for image_path in image_paths]
+                            for _ in range(request.num_completions)
+                        ],
+                        return_tensors="pt",
+                    )
+                    inputs = {k: v.to(self._device) for k, v in inputs.items()}
+
+                    # Generate
+                    generated_ids = model.generate(**inputs, **generation_args)  # type: ignore
+                    generated_texts: List[str] = processor.batch_decode(  # type: ignore
+                        generated_ids, skip_special_tokens=True
+                    )
+                    return {"output": generated_texts}
+
+                # Include the prompt and model name in the cache key
+                cache_key = CachingClient.make_cache_key(
+                    raw_request={
+                        "n": request.num_completions,
+                        "model": request.model,
+                        "prompt": generate_uid_for_multimodal_prompt(request.multimodal_prompt),
+                        **generation_args,
+                    },
+                    request=request,
+                )
+                result, cached = self.cache.get(cache_key, wrap_request_time(do_it))
+            except RuntimeError as model_error:
+                hexception(model_error)
+                return RequestResult(success=False, cached=False, error=str(model_error), completions=[], embedding=[])
+
+            for text in result["output"]:
+                hlog(f"Generated text: {text}")
+                assert self.ASSISTANT_PREFIX in text, f"Expected {self.ASSISTANT_PREFIX} in the output"
+                text = text.rpartition(self.ASSISTANT_PREFIX)[-1]
+                hlog(f"Truncated: {text}")
+                tokenization_result = self.tokenizer.tokenize(
+                    TokenizationRequest(text, tokenizer=self.tokenizer_name, encode=False)
+                )
+                tokens: List[Token] = [Token(text=str(text), logprob=0) for text in tokenization_result.raw_tokens]
+                completions.append(GeneratedOutput(text=text, logprob=0, tokens=tokens))
+
+        return RequestResult(
+            success=True,
+            cached=cached,
+            request_time=result["request_time"],
+            completions=completions,
+            embedding=[],
+        )