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

helm/clients/vertexai_client.py

@@ -1,10 +1,9 @@
 import requests
 from abc import ABC, abstractmethod
 from threading import Lock
-from typing import Any, Dict, Optional, List, Union
+from typing import Any, Dict, Mapping, Optional, List, Union
 
 from helm.common.cache import CacheConfig
-from helm.common.hierarchical_logger import hlog
 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
@@ -27,22 +26,62 @@ 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}")
+
+
+def _get_model_name_for_request(request: Request) -> str:
+    # We have to strip "-safety-" suffixes from model names because they are not part of the Vertex AI model name
+    # TODO: Clean up this hack
+    return request.model_engine.split("-safety-")[0]
+
+
 class VertexAIClient(CachingClient, ABC):
     """Client for Vertex AI models"""
 
-    def __init__(self, cache_config: CacheConfig, project_id: str, location: str) -> None:
+    def __init__(
+        self, cache_config: CacheConfig, project_id: str, location: str, safety_settings_preset: Optional[str] = None
+    ) -> None:
         super().__init__(cache_config=cache_config)
         self.project_id = project_id
         self.location = location
 
-        # VertexAI's default safety filter is overly sensitive, so we disable it.
-        self.safety_settings: Dict[HarmCategory, SafetySetting.HarmBlockThreshold] = {
-            harm_category: SafetySetting.HarmBlockThreshold(SafetySetting.HarmBlockThreshold.BLOCK_NONE)
-            for harm_category in iter(HarmCategory)
-        }
+        self.safety_settings_preset = safety_settings_preset
+        self.safety_settings = _get_safety_settings_for_preset(safety_settings_preset)
 
         vertexai.init(project=self.project_id, location=self.location)
 
+    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
@@ -72,7 +111,7 @@ class VertexAITextClient(VertexAIClient):
         }
 
         completions: List[GeneratedOutput] = []
-        model_name: str = request.model_engine
+        model_name: str = _get_model_name_for_request(request)
 
         try:
 
@@ -88,9 +127,9 @@ class VertexAITextClient(VertexAIClient):
             # 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 = CachingClient.make_cache_key(
+            cache_key = self.make_cache_key_with_safety_settings_preset(
                 {
-                    "engine": request.model_engine,
+                    "engine": model_name,
                     "prompt": request.prompt,
                     **parameters,
                 },
@@ -131,12 +170,6 @@ class VertexAITextClient(VertexAIClient):
 class VertexAIChatClient(VertexAIClient):
     """Client for Vertex AI chat models (e.g., Gemini). Supports multimodal prompts."""
 
-    # Set the finish reason to this if the prompt violates the content policy
-    CONTENT_POLICY_VIOLATED_FINISH_REASON: str = "The prompt violates Google's content policy."
-
-    # Gemini returns this error for certain valid requests
-    CONTENT_HAS_NO_PARTS_ERROR: str = "Content has no parts."
-
     # 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.
@@ -149,7 +182,7 @@ class VertexAIChatClient(VertexAIClient):
     ]
 
     @staticmethod
-    def get_model(model_name: str) -> Any:
+    def get_model(model_name: str) -> GenerativeModel:
         global _models_lock
         global _models
 
@@ -184,7 +217,7 @@ class VertexAIChatClient(VertexAIClient):
         }
 
         completions: List[GeneratedOutput] = []
-        model_name: str = request.model_engine
+        model_name: str = _get_model_name_for_request(request)
         model = self.get_model(model_name)
 
         try:
@@ -202,21 +235,24 @@ class VertexAIChatClient(VertexAIClient):
                 )
                 candidates: List[Candidate] = response.candidates
 
-                # Depending on the version of the Vertex AI library and the type of content blocking,
-                # content blocking can show up in many ways, so this defensively handles most of these ways
+                # 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("No candidates in response due to content blocking")
+                    raise VertexAIContentBlockedError(f"No candidates in response: {response}")
                 predictions: List[Dict[str, Any]] = []
                 for candidate in candidates:
-                    if (
-                        candidate.finish_reason in VertexAIChatClient.CONTENT_BLOCKED_FINISH_REASONS
-                        or not candidate.content.parts
-                    ):
-                        # The prediction was either blocked due to safety settings or the model stopped and returned
-                        # nothing (which also happens when the model is blocked).
-                        # For now, we don't cache blocked requests, because we are trying to get the
-                        # content blocking removed.
-                        raise VertexAIContentBlockedError("Content has no parts due to content blocking")
+                    # 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}")
                     predictions.append({"text": candidate.content.text})
                     # TODO: Extract more information from the response
                 return {"predictions": predictions}
@@ -224,7 +260,7 @@ class VertexAIChatClient(VertexAIClient):
             # 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 = CachingClient.make_cache_key(
+            cache_key = self.make_cache_key_with_safety_settings_preset(
                 {
                     "model_name": model_name,
                     "prompt": request.prompt,
@@ -234,11 +270,11 @@ class VertexAIChatClient(VertexAIClient):
             )
 
             response, cached = self.cache.get(cache_key, wrap_request_time(do_it))
-        except VertexAIContentBlockedError:
+        except VertexAIContentBlockedError as e:
             return RequestResult(
                 success=False,
                 cached=False,
-                error="Response was empty due to content moderation filter",
+                error=f"Content blocked: {str(e)}",
                 completions=[],
                 embedding=[],
                 error_flags=ErrorFlags(is_retriable=False, is_fatal=False),
@@ -252,7 +288,7 @@ class VertexAIChatClient(VertexAIClient):
             return RequestResult(
                 success=False,
                 cached=False,
-                error="Response was empty due to content moderation filter",
+                error=f"Content blocked error in cached response: {str(response)}",
                 completions=[],
                 embedding=[],
                 error_flags=ErrorFlags(is_retriable=False, is_fatal=False),
@@ -266,7 +302,7 @@ class VertexAIChatClient(VertexAIClient):
                 return RequestResult(
                     success=False,
                     cached=False,
-                    error="Response was empty due to content moderation filter",
+                    error=f"Content blocked error in cached prediction: {str(prediction)}",
                     completions=[],
                     embedding=[],
                     error_flags=ErrorFlags(is_retriable=False, is_fatal=False),
@@ -291,21 +327,6 @@ class VertexAIChatClient(VertexAIClient):
         )
 
     def _make_multimodal_request(self, request: Request) -> RequestResult:
-        def complete_for_valid_error(error_message: str) -> RequestResult:
-            empty_completion = GeneratedOutput(
-                text="",
-                logprob=0,
-                tokens=[],
-                finish_reason={"reason": error_message},
-            )
-            return RequestResult(
-                success=True,
-                cached=False,
-                request_time=0,
-                completions=[empty_completion] * request.num_completions,
-                embedding=[],
-            )
-
         # 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
@@ -334,7 +355,7 @@ class VertexAIChatClient(VertexAIClient):
         }
 
         completions: List[GeneratedOutput] = []
-        model_name: str = request.model_engine
+        model_name: str = _get_model_name_for_request(request)
         model = self.get_model(model_name)
 
         request_time = 0
@@ -346,30 +367,62 @@ class VertexAIChatClient(VertexAIClient):
             try:
 
                 def do_it() -> Dict[str, Any]:
-                    raw_response = model.generate_content(
+                    response: GenerationResponse = model.generate_content(
                         contents, generation_config=parameters, safety_settings=self.safety_settings
                     )
-                    if raw_response._raw_response.prompt_feedback.block_reason != 0:
-                        hlog(f"Content blocked for prompt: {request.multimodal_prompt}")
-                        return {"error": self.CONTENT_POLICY_VIOLATED_FINISH_REASON}
-
-                    return {"predictions": [{"text": raw_response.candidates[0].text}]}
+                    # 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 = CachingClient.make_cache_key(raw_cache_key, request)
+                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, ValueError) as e:
-                if str(e) == self.CONTENT_HAS_NO_PARTS_ERROR:
-                    return complete_for_valid_error(self.CONTENT_HAS_NO_PARTS_ERROR)
-
+            except requests.exceptions.RequestException as 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 complete_for_valid_error(response["error"])
+                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=[])

helm/clients/vision_language/huggingface_vision2seq_client.py

@@ -0,0 +1,145 @@
+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 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)
+                    inputs = processor(
+                        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)
+                    generated_texts: List[str] = processor.batch_decode(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:
+                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=[],
+        )

helm/clients/vision_language/huggingface_vlm_client.py

@@ -25,7 +25,7 @@ except ModuleNotFoundError as e:
 
 class HuggingFaceVLMClient(CachingClient):
     """
-    General CLient for VLM models from HuggingFace.
+    General client for VLM models from HuggingFace.
     """
 
     _models_lock: Lock = Lock()
@@ -34,6 +34,11 @@ class HuggingFaceVLMClient(CachingClient):
         "huggingface/llava-1.5-7b-hf": "llava-hf/llava-1.5-7b-hf",
         "huggingface/llava-1.5-13b-hf": "llava-hf/llava-1.5-13b-hf",
         "huggingface/bakLlava-v1-hf": "llava-hf/bakLlava-v1-hf",
+        "huggingface/llava-v1.6-vicuna-7b-hf": "llava-hf/llava-v1.6-vicuna-7b-hf",
+        "huggingface/llava-v1.6-vicuna-13b-hf": "llava-hf/llava-v1.6-vicuna-13b-hf",
+        "huggingface/llava-v1.6-mistral-7b-hf": "llava-hf/llava-v1.6-mistral-7b-hf",
+        "huggingface/llava-v1.6-34b-hf": "llava-hf/llava-v1.6-34b-hf",
+        "huggingface/prometheus-vision-13b-v1.0-hf": "PahaII/prometheus-vision-13b-v1.0-hf",
     }
 
     def __init__(self, tokenizer: Tokenizer, tokenizer_name: str, cache_config: CacheConfig):
@@ -45,7 +50,7 @@ class HuggingFaceVLMClient(CachingClient):
         with self._models_lock:
             model_id: str = self._models_aliases.get(model_name, model_name)
             if model_id not in self._models:
-                self._models[model_id] = pipeline("image-to-text", model=model_id)
+                self._models[model_id] = pipeline("image-to-text", model=model_id, device_map="auto")
             return self._models[model_id]
 
     def make_request(self, request: Request) -> RequestResult:
@@ -90,11 +95,14 @@ class HuggingFaceVLMClient(CachingClient):
         except RuntimeError as e:
             return RequestResult(success=False, cached=False, error=str(e), completions=[], embedding=[])
 
+        output: str = result["generated_text"]
+        if "ASSISTANT: " in output:
+            output = output.split("ASSISTANT: ")[1]
         tokenization_result: TokenizationRequestResult = self.tokenizer.tokenize(
-            TokenizationRequest(result["generated_text"], tokenizer=self.tokenizer_name)
+            TokenizationRequest(output, tokenizer=self.tokenizer_name)
         )
         tokens: List[Token] = [Token(text=str(text), logprob=0) for text in tokenization_result.raw_tokens]
-        completions: List[GeneratedOutput] = [GeneratedOutput(text=result["generated_text"], logprob=0, tokens=tokens)]
+        completions: List[GeneratedOutput] = [GeneratedOutput(text=output, logprob=0, tokens=tokens)]
         return RequestResult(
             success=True,
             cached=cached,

helm/clients/vision_language/idefics_client.py

@@ -88,7 +88,7 @@ class IDEFICSClient(CachingClient):
 
         input_args: Dict[str, Union[str, bool]] = {"return_tensors": "pt"}
         generation_args = {
-            "max_length": request.max_tokens,
+            "max_new_tokens": request.max_tokens,
             "bad_words_ids": processor.tokenizer(self.BAD_WORD_TOKENS, add_special_tokens=False).input_ids,
         }
 
@@ -140,7 +140,7 @@ class IDEFICSClient(CachingClient):
 
                 # Truncate the output text as IDEFICS outputs the entire sequence including the prompt
                 if "instruct" in request.model:
-                    assert self.ASSISTANT_PREFIX in text, f"Expected {self.ASSISTANT_PREFIX} in the output"
+                    assert self.ASSISTANT_PREFIX in text, f"Expected {self.ASSISTANT_PREFIX} in the output: {text}"
                     text = text.rpartition(self.ASSISTANT_PREFIX)[-1]
                 else:
                     # Best we can do is to remove the text portion of the prompt from the output

helm/clients/vision_language/paligemma_client.py

@@ -0,0 +1,146 @@
+from threading import Lock
+from typing import Any, Dict, List, Optional
+
+import torch
+from dataclasses import dataclass
+from transformers import AutoProcessor, PaliGemmaForConditionalGeneration
+
+from helm.common.cache import CacheConfig
+from helm.common.images_utils import open_image
+from helm.common.gpu_utils import get_torch_device_name
+from helm.common.hierarchical_logger import hlog, htrack_block
+from helm.common.media_object import TEXT_TYPE
+from helm.common.optional_dependencies import handle_module_not_found_error
+from helm.common.request import Request, RequestResult, GeneratedOutput, Token
+from helm.common.tokenization_request import TokenizationRequest
+from helm.common.request import wrap_request_time
+from helm.clients.client import CachingClient, generate_uid_for_multimodal_prompt
+from helm.tokenizers.tokenizer import Tokenizer
+
+try:
+    from PIL import Image
+except ModuleNotFoundError as e:
+    handle_module_not_found_error(e, ["images"])
+
+# Added to solve: cutlassF: no kernel found to launch!
+torch.backends.cuda.enable_mem_efficient_sdp(False)
+torch.backends.cuda.enable_flash_sdp(False)
+
+
+@dataclass(frozen=True)
+class LoadedPaliGemmaForConditionalGeneration:
+    """Loaded model and processor for PaliGemma."""
+
+    model: PaliGemmaForConditionalGeneration
+    processor: AutoProcessor
+
+
+_models_lock: Lock = Lock()
+_models: Dict[str, Optional[LoadedPaliGemmaForConditionalGeneration]] = {}
+
+
+class PaliGemmaClient(CachingClient):
+    """
+    PaliGemma is a versatile and lightweight vision-language model (VLM) inspired by PaLI-3
+    and based on open components such as the SigLIP vision model and the Gemma language model.
+    It takes both image and text as input and generates text as output, supporting multiple languages.
+    It is designed for class-leading fine-tune performance on a wide range of vision-language tasks
+    such as image and short video caption, visual question answering, text reading, object detection
+    and object segmentation.
+    """
+
+    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) -> LoadedPaliGemmaForConditionalGeneration:
+        global _models_lock
+        global _models
+
+        # Ensure that only one thread is loading the model at a time
+        with _models_lock:
+            if checkpoint not in _models or _models[checkpoint] is None:
+                hlog(f"Loading model {checkpoint} and caching in memory...")
+                model = PaliGemmaForConditionalGeneration.from_pretrained(
+                    checkpoint, torch_dtype=torch.bfloat16, device_map="auto"
+                ).eval()
+                processor = AutoProcessor.from_pretrained(checkpoint)
+                _models[checkpoint] = LoadedPaliGemmaForConditionalGeneration(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.multimodal_prompt is not None, "Multimodal prompt is required"
+
+        loaded_model_processor: LoadedPaliGemmaForConditionalGeneration = self._get_model(request.model_deployment)
+        model = loaded_model_processor.model
+        processor = loaded_model_processor.processor
+        generation_args = {"max_new_tokens": request.max_tokens}
+
+        images: List[Image.Image] = []
+        prompt_pieces: List[str] = []
+        for media_object in request.multimodal_prompt.media_objects:
+            if media_object.is_type("image") and media_object.location:
+                images += [open_image(media_object.location).convert("RGB")]
+            elif media_object.is_type(TEXT_TYPE):
+                if media_object.text is None:
+                    raise ValueError("MediaObject of text type has missing text field value")
+                prompt_pieces.append(media_object.text)
+            else:
+                raise ValueError(f"Unrecognized MediaObject type {media_object.type}")
+        prompt_text: str = "\n".join(prompt_pieces)
+        model_inputs = processor(text=prompt_text, images=images, return_tensors="pt").to(self._device)
+        input_len = model_inputs["input_ids"].shape[-1]
+
+        completions: List[GeneratedOutput] = []
+        with htrack_block(f"Generating for prompt: {prompt_text}"):
+            try:
+                concat_results = []
+                for i_completion in range(request.num_completions):
+
+                    def do_it() -> Dict[str, Any]:
+                        with torch.inference_mode():
+                            generation = model.generate(
+                                **model_inputs, max_new_tokens=request.max_tokens, do_sample=False
+                            )[0]
+                            if not request.echo_prompt:
+                                generation = generation[input_len:]
+                            decoded = processor.decode(generation, skip_special_tokens=True)
+                            return {"output": decoded}
+
+                    # Include the prompt and model name in the cache key
+                    cache_key = CachingClient.make_cache_key(
+                        raw_request={
+                            "n": request.num_completions,
+                            "i": i_completion,
+                            "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))
+                    concat_results.append(result)
+            except RuntimeError as model_error:
+                return RequestResult(success=False, cached=False, error=str(model_error), completions=[], embedding=[])
+
+            for result in concat_results:
+                text = result["output"]
+                hlog(f"Generated text: {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=[],
+        )