crfm-helm 0.5.4__py3-none-any.whl → 0.5.6__py3-none-any.whl

helm/clients/vertexai_client.py

@@ -4,6 +4,7 @@ from threading import Lock
 from typing import Any, Dict, Mapping, Optional, List, Union
 
 from helm.common.cache import CacheConfig
+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
@@ -12,7 +13,14 @@ from helm.clients.client import CachingClient, truncate_sequence, generate_uid_f
 try:
     import vertexai
     from vertexai.language_models import TextGenerationModel, TextGenerationResponse  # PaLM2
-    from vertexai.preview.generative_models import GenerativeModel, GenerationResponse, Candidate, Part, Image  # Gemini
+    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"])
@@ -48,17 +56,16 @@ def _get_safety_settings_for_preset(
         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, safety_settings_preset: Optional[str] = None
+        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
@@ -67,8 +74,15 @@ class VertexAIClient(CachingClient, ABC):
         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.
 
@@ -111,7 +125,7 @@ class VertexAITextClient(VertexAIClient):
         }
 
         completions: List[GeneratedOutput] = []
-        model_name: str = _get_model_name_for_request(request)
+        model_name: str = self._get_model_name_for_request(request)
 
         try:
 
@@ -193,12 +207,20 @@ class VertexAIChatClient(VertexAIClient):
 
     def make_request(self, request: Request) -> RequestResult:
         """Make a request"""
-        contents: str = request.prompt
+        contents = [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:
+            contents = []
+            role_mapping = {"user": "user", "assistant": "model"}
+            for msg in request.messages:
+                contents.append(
+                    Content(role=role_mapping.get(msg["role"], "user"), parts=[Part.from_text(msg["content"])])
+                )
+
         parameters = {
             "temperature": request.temperature,
             "max_output_tokens": request.max_tokens,
@@ -217,7 +239,7 @@ class VertexAIChatClient(VertexAIClient):
         }
 
         completions: List[GeneratedOutput] = []
-        model_name: str = _get_model_name_for_request(request)
+        model_name: str = self._get_model_name_for_request(request)
         model = self.get_model(model_name)
 
         try:
@@ -263,7 +285,7 @@ class VertexAIChatClient(VertexAIClient):
             cache_key = self.make_cache_key_with_safety_settings_preset(
                 {
                     "model_name": model_name,
-                    "prompt": request.prompt,
+                    "prompt": request.messages or request.prompt,
                     **parameters,
                 },
                 request,
@@ -338,6 +360,16 @@ class VertexAIChatClient(VertexAIClient):
         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")
@@ -355,7 +387,7 @@ class VertexAIChatClient(VertexAIClient):
         }
 
         completions: List[GeneratedOutput] = []
-        model_name: str = _get_model_name_for_request(request)
+        model_name: str = self._get_model_name_for_request(request)
         model = self.get_model(model_name)
 
         request_time = 0

helm/clients/vision_language/huggingface_vision2seq_client.py

@@ -95,8 +95,8 @@ class HuggingFaceVision2SeqClient(CachingClient):
 
                 def do_it() -> Dict[str, Any]:
                     messages = [{"role": "user", "content": multimodal_prompt}]
-                    prompt = processor.apply_chat_template(messages, add_generation_prompt=True)
-                    inputs = processor(
+                    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]
@@ -107,8 +107,10 @@ class HuggingFaceVision2SeqClient(CachingClient):
                     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)
+                    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

helm/clients/vision_language/huggingface_vlm_client.py

@@ -50,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, device_map="auto")
+                self._models[model_id] = pipeline("image-to-text", model=model_id, device_map="auto")  # type: ignore
             return self._models[model_id]
 
     def make_request(self, request: Request) -> RequestResult:
@@ -80,7 +80,7 @@ class HuggingFaceVLMClient(CachingClient):
 
             def do_it() -> Dict[str, Any]:
                 model: ImageToTextPipeline = self._get_model(request.model_deployment)
-                outputs = model(image, prompt=prompt, generate_kwargs=generation_args)
+                outputs = model(image, prompt=prompt, generate_kwargs=generation_args)  # type: ignore
                 return outputs[0]
 
             cache_key = CachingClient.make_cache_key(

helm/clients/vision_language/idefics_client.py

@@ -89,14 +89,18 @@ class IDEFICSClient(CachingClient):
         input_args: Dict[str, Union[str, bool]] = {"return_tensors": "pt"}
         generation_args = {
             "max_new_tokens": request.max_tokens,
-            "bad_words_ids": processor.tokenizer(self.BAD_WORD_TOKENS, add_special_tokens=False).input_ids,
+            "bad_words_ids": processor.tokenizer(  # type: ignore
+                self.BAD_WORD_TOKENS, add_special_tokens=False
+            ).input_ids,
         }
 
         if self.END_OF_UTTERANCE_TOKEN in request.stop_sequences:
             # Following https://huggingface.co/HuggingFaceM4/idefics-80b-instruct,
             # specify <end_of_utterance> as an exit condition.
             input_args["add_end_of_utterance_token"] = False
-            exit_condition = processor.tokenizer(self.END_OF_UTTERANCE_TOKEN, add_special_tokens=False).input_ids
+            exit_condition = processor.tokenizer(  # type: ignore
+                self.END_OF_UTTERANCE_TOKEN, add_special_tokens=False
+            ).input_ids
             generation_args["eos_token_id"] = exit_condition
 
         multimodal_prompt: List[Union[str, Image.Image]] = []

helm/clients/vision_language/open_flamingo/__init__.py

@@ -1,2 +1,2 @@
-from .src.flamingo import Flamingo
-from .src.factory import create_model_and_transforms
+from helm.clients.vision_language.open_flamingo.src.flamingo import Flamingo
+from helm.clients.vision_language.open_flamingo.src.factory import create_model_and_transforms

helm/clients/vision_language/open_flamingo/src/factory.py

@@ -7,9 +7,9 @@ from typing import Optional
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 from helm.common.general import handle_module_not_found_error
-from .flamingo import Flamingo
-from .flamingo_lm import FlamingoLMMixin
-from .utils import extend_instance
+from helm.clients.vision_language.open_flamingo.src.flamingo import Flamingo
+from helm.clients.vision_language.open_flamingo.src.flamingo_lm import FlamingoLMMixin
+from helm.clients.vision_language.open_flamingo.src.utils import extend_instance
 
 
 def create_model_and_transforms(

helm/clients/vision_language/open_flamingo/src/flamingo.py

@@ -5,7 +5,7 @@ Source: https://github.com/mlfoundations/open_flamingo
 import torch
 from einops import rearrange
 from torch import nn
-from .helpers import PerceiverResampler
+from helm.clients.vision_language.open_flamingo.src.helpers import PerceiverResampler
 from torch.distributed.fsdp.wrap import (
     enable_wrap,
     wrap,
@@ -15,7 +15,7 @@ from torch.distributed.fsdp import (
     FullyShardedDataParallel as FSDP,
 )
 
-from .utils import apply_with_stopping_condition
+from helm.clients.vision_language.open_flamingo.src.utils import apply_with_stopping_condition
 
 
 class Flamingo(nn.Module):

helm/clients/vision_language/open_flamingo/src/flamingo_lm.py

@@ -3,8 +3,8 @@ Source: https://github.com/mlfoundations/open_flamingo
 """
 
 import torch.nn as nn
-from .helpers import GatedCrossAttentionBlock
-from .utils import getattr_recursive, setattr_recursive
+from helm.clients.vision_language.open_flamingo.src.helpers import GatedCrossAttentionBlock
+from helm.clients.vision_language.open_flamingo.src.utils import getattr_recursive, setattr_recursive
 
 
 class FlamingoLayer(nn.Module):

helm/clients/vision_language/paligemma_client.py

@@ -93,7 +93,7 @@ class PaliGemmaClient(CachingClient):
             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)
+        model_inputs = processor(text=prompt_text, images=images, return_tensors="pt").to(self._device)  # type: ignore
         input_len = model_inputs["input_ids"].shape[-1]
 
         completions: List[GeneratedOutput] = []
@@ -109,7 +109,7 @@ class PaliGemmaClient(CachingClient):
                             )[0]
                             if not request.echo_prompt:
                                 generation = generation[input_len:]
-                            decoded = processor.decode(generation, skip_special_tokens=True)
+                            decoded = processor.decode(generation, skip_special_tokens=True)  # type: ignore
                             return {"output": decoded}
 
                     # Include the prompt and model name in the cache key

helm/clients/vision_language/qwen2_vlm_client.py

@@ -0,0 +1,188 @@
+from threading import Lock
+from typing import Any, Dict, List, Optional
+from dataclasses import dataclass
+
+from transformers import AutoProcessor
+from qwen_vl_utils import process_vision_info
+import torch
+
+from helm.common.cache import CacheConfig
+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.request import Request, RequestResult, GeneratedOutput, Token
+from helm.common.request import wrap_request_time
+from helm.clients.client import CachingClient, generate_uid_for_multimodal_prompt
+
+
+@dataclass(frozen=True)
+class LoadedModelProcessor:
+    model: Any
+    processor: AutoProcessor
+
+
+# Global cache for all models
+_models_lock: Lock = Lock()
+_models: Dict[str, Optional[LoadedModelProcessor]] = {
+    "Qwen/Qwen2-VL-7B-Instruct": None,
+    "Qwen/Qwen2-VL-72B-Instruct": None,
+    "Qwen/Qwen2.5-VL-3B-Instruct": None,
+    "Qwen/Qwen2.5-VL-7B-Instruct": None,
+    "Qwen/Qwen2.5-VL-32B-Instruct": None,
+    "Qwen/Qwen2.5-VL-72B-Instruct": None,
+}
+
+
+class Qwen2VLMClient(CachingClient):
+    def __init__(self, cache_config: CacheConfig):
+        super().__init__(cache_config=cache_config)
+        self._device: str = get_torch_device_name()
+
+    def _get_model_name(self, helm_model_name: str) -> str:
+        if helm_model_name == "qwen2-vl-7b-instruct":
+            return "Qwen/Qwen2-VL-7B-Instruct"
+        elif helm_model_name == "qwen2-vl-72b-instruct":
+            return "Qwen/Qwen2-VL-72B-Instruct"
+        elif helm_model_name == "qwen2.5-vl-3b-instruct":
+            return "Qwen/Qwen2.5-VL-3B-Instruct"
+        elif helm_model_name == "qwen2.5-vl-7b-instruct":
+            return "Qwen/Qwen2.5-VL-7B-Instruct"
+        elif helm_model_name == "qwen2.5-vl-32b-instruct":
+            return "Qwen/Qwen2.5-VL-32B-Instruct"
+        elif helm_model_name == "qwen2.5-vl-72b-instruct":
+            return "Qwen/Qwen2.5-VL-72B-Instruct"
+        else:
+            raise ValueError(f"Unhandled model name: {helm_model_name}")
+
+    def _get_model(self, helm_model_name: str) -> LoadedModelProcessor:
+        from transformers import Qwen2VLForConditionalGeneration, Qwen2_5_VLForConditionalGeneration
+
+        global _models_lock, _models
+
+        model_name = self._get_model_name(helm_model_name)
+        with _models_lock:
+            loaded = _models[model_name]
+            if loaded is None:
+                hlog(f"Loading model {model_name} and caching in memory...")
+                # Use different loading routines depending on whether it's Qwen2.5 or Qwen2.
+                if "2.5" in model_name:
+                    # Qwen2.5: by default use torch_dtype="auto". You can enable flash_attention_2 if desired.
+                    model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+                        model_name,
+                        torch_dtype=torch.bfloat16,
+                        device_map="auto",
+                        attn_implementation="flash_attention_2",
+                    ).eval()
+                else:
+                    model = Qwen2VLForConditionalGeneration.from_pretrained(
+                        model_name,
+                        torch_dtype=torch.bfloat16,
+                        device_map="auto",
+                        attn_implementation="flash_attention_2",
+                    ).eval()
+                processor = AutoProcessor.from_pretrained(model_name)
+                loaded = LoadedModelProcessor(model=model, processor=processor)
+                _models[model_name] = loaded
+        return loaded
+
+    def make_request(self, request: Request) -> RequestResult:
+        assert request.multimodal_prompt is not None, "Multimodal prompt is required"
+
+        # Build messages by collating all media objects into a single "user" message.
+        message_content = []
+        for media_object in request.multimodal_prompt.media_objects:
+            if media_object.is_type("image") and media_object.location:
+                message_content.append({"type": "image", "image": media_object.location})
+            elif media_object.is_type(TEXT_TYPE):
+                if media_object.text is None:
+                    raise ValueError("MediaObject of text type has missing text field value")
+                message_content.append({"type": "text", "text": media_object.text})
+            else:
+                raise ValueError(f"Unrecognized MediaObject type {media_object.type}")
+
+        messages = [{"role": "user", "content": message_content}]
+
+        generation_args = {
+            "max_new_tokens": request.max_tokens,
+        }
+
+        completions: List[GeneratedOutput] = []
+        request_time: float = 0
+        request_datetime: Optional[int] = None
+        all_cached: bool = True
+
+        with htrack_block(f"Generating for prompt: {request.multimodal_prompt.text}"):
+            for completion_index in range(request.num_completions):
+                try:
+
+                    def do_it() -> Dict[str, Any]:
+                        loaded = self._get_model(request.model_engine)
+                        model = loaded.model
+                        processor = loaded.processor
+
+                        # Prepare text and vision inputs.
+                        text = processor.apply_chat_template(  # type: ignore
+                            messages, tokenize=False, add_generation_prompt=True
+                        )
+                        image_inputs, video_inputs = process_vision_info(messages)
+                        inputs = processor(  # type: ignore
+                            text=[text],
+                            images=image_inputs,
+                            videos=video_inputs,
+                            padding=True,
+                            return_tensors="pt",
+                        ).to(self._device)
+
+                        generated_ids = model.generate(**inputs, **generation_args)
+                        # Remove the input prefix from outputs.
+                        generated_ids_trimmed = [
+                            out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+                        ]
+                        output_text = processor.batch_decode(  # type: ignore
+                            generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
+                        )
+                        # For simplicity, we split tokens by whitespace.
+                        completion = output_text[0]
+                        tokens = completion.split()
+                        return {"output": (completion, tokens)}
+
+                    cache_key = CachingClient.make_cache_key(
+                        raw_request={
+                            "completion_index": completion_index,
+                            "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=[],
+                    )
+
+                text_out, tokens = result["output"]
+                completions.append(
+                    GeneratedOutput(
+                        text=text_out,
+                        logprob=0,
+                        tokens=[Token(text=str(token), logprob=0) for token in tokens],
+                    )
+                )
+                hlog(f"Generated: {text_out}")
+                request_time += result["request_time"]
+                request_datetime = request_datetime or result.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/qwen_vlm_client.py

@@ -115,14 +115,16 @@ class QwenVLMClient(CachingClient):
 
                     def do_it() -> Dict[str, Any]:
                         if request.model_engine == "qwen-vl-chat":
-                            completion, _ = model.chat(tokenizer, query=tokenizer.from_list_format(query), history=None)
+                            completion, _ = model.chat(  # type: ignore
+                                tokenizer, query=tokenizer.from_list_format(query), history=None  # type: ignore
+                            )
                         else:
-                            inputs = tokenizer(tokenizer.from_list_format(query), return_tensors="pt")
+                            inputs = tokenizer(tokenizer.from_list_format(query), return_tensors="pt")  # type: ignore
                             inputs = inputs.to(self._device)
-                            pred = model.generate(**inputs, **generation_args)
-                            completion = tokenizer.decode(pred.cpu()[0], skip_special_tokens=False)
+                            pred = model.generate(**inputs, **generation_args)  # type: ignore
+                            completion = tokenizer.decode(pred.cpu()[0], skip_special_tokens=False)  # type: ignore
 
-                        tokens: List[str] = tokenizer.tokenize(completion)
+                        tokens: List[str] = tokenizer.tokenize(completion)  # type: ignore
                         return {"output": (completion, tokens)}
 
                     # Include the prompt and model name in the cache key

helm/clients/vllm_client.py

@@ -2,13 +2,15 @@ from typing import Any, Dict, Optional
 
 from helm.common.cache import CacheConfig
 from helm.common.request import Request
-from helm.clients.openai_client import OpenAIClient
+from helm.clients.openai_client import OpenAILegacyCompletionsClient
 from helm.tokenizers.tokenizer import Tokenizer
 
 
-class VLLMClient(OpenAIClient):
+class VLLMClient(OpenAILegacyCompletionsClient):
     """Sends request to a vLLM server using the OpenAI-compatible API.
 
+    Only supports the legacy Text Completions API, rather than the Chat Completions API.
+
     See: https://docs.vllm.ai/en/latest/getting_started/quickstart.html#openai-compatible-server"""
 
     def __init__(
@@ -29,10 +31,6 @@ class VLLMClient(OpenAIClient):
         self.tokenizer = tokenizer
         self.tokenizer_name = tokenizer_name
 
-    def _is_chat_model_engine(self, model_engine: str) -> bool:
-        # Only support vLLM completion models for now.
-        return False
-
     def _get_model_for_request(self, request: Request) -> str:
         # The `model` parameter for vLLM should be the whole model name including the creator organization,
         # unlike OpenAI which only uses the model engine.

helm/clients/writer_client.py

@@ -0,0 +1,102 @@
+from typing import Any, Dict, List, Mapping, Optional
+
+from helm.clients.client import CachingClient
+from helm.common.cache import CacheConfig
+from helm.common.optional_dependencies import handle_module_not_found_error
+from helm.common.request import wrap_request_time, Request, RequestResult, GeneratedOutput, Token
+
+try:
+    from writerai import Writer
+    from writerai.types.chat_completion import ChatCompletion
+except ModuleNotFoundError as e:
+    handle_module_not_found_error(e, ["openai"])
+
+
+class WriterClient(CachingClient):
+    def __init__(self, cache_config: CacheConfig, api_key: Optional[str] = None):
+        super().__init__(cache_config=cache_config)
+        self._writer_client = Writer(api_key=api_key)
+
+    def _get_messages_from_request(self, request: Request) -> List[Dict]:
+        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("`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_chat_completion_to_generated_outputs(
+        self, chat_completion: ChatCompletion, request: Request
+    ) -> List[GeneratedOutput]:
+        generated_outputs: List[GeneratedOutput] = []
+        for choice in chat_completion.choices:
+            raw_completion_content = choice.message.content
+            # The Writer chat completion API doesn't support echo.
+            # If `echo_prompt` is true, combine the prompt and completion.
+            text: str = request.prompt + raw_completion_content if request.echo_prompt else raw_completion_content
+            tokens: List[Token] = []
+            if choice.logprobs and choice.logprobs.content:
+                tokens = [
+                    Token(text=choice_token.token, logprob=choice_token.logprob)
+                    for choice_token in choice.logprobs.content
+                ]
+            generated_output = GeneratedOutput(
+                text=text,
+                logprob=sum(token.logprob for token in tokens) if tokens else 0.0,
+                tokens=tokens,
+                finish_reason={"reason": choice.finish_reason},
+            )
+            generated_outputs.append(generated_output)
+        return generated_outputs
+
+    def _convert_request_to_raw_request(self, request: Request) -> Dict:
+        raw_request = {
+            "messages": self._get_messages_from_request(request),
+            "model": request.model.split("/")[-1],
+            "logprobs": bool(request.top_k_per_token),
+            "max_tokens": request.max_tokens,
+            "n": request.num_completions,
+            "stop": request.stop_sequences,
+            "temperature": request.temperature,
+            "top_p": request.top_p,
+        }
+        if request.response_format and request.response_format.json_schema:
+            raw_request["response_format"] = {
+                "type": "json_schema",
+                "json_schema": {
+                    "schema": request.response_format.json_schema,
+                },
+            }
+        return raw_request
+
+    def make_request(self, request: Request) -> RequestResult:
+        raw_request = self._convert_request_to_raw_request(request)
+        cache_key: Mapping = CachingClient.make_cache_key(raw_request, request)
+
+        def do_it() -> Dict[Any, Any]:
+            return self._writer_client.chat.chat(**raw_request).model_dump()
+
+        try:
+            raw_response, cached = self.cache.get(cache_key, wrap_request_time(do_it))
+            chat_completion: ChatCompletion = ChatCompletion.model_validate(raw_response)
+        except Exception as error:
+            return RequestResult(
+                success=False,
+                cached=False,
+                error=str(error),
+                completions=[],
+                embedding=[],
+            )
+
+        generated_outputs = self._convert_chat_completion_to_generated_outputs(chat_completion, request)
+
+        return RequestResult(
+            success=True,
+            cached=cached,
+            request_time=raw_response["request_time"],
+            request_datetime=raw_response["request_datetime"],
+            completions=generated_outputs,
+            embedding=[],
+        )

helm/clients/yi_client.py

@@ -26,6 +26,3 @@ class YiChatClient(OpenAIClient):
             org_id=None,
             base_url=YiChatClient.BASE_URL,
         )
-
-    def _is_chat_model_engine(self, model_engine: str) -> bool:
-        return True