[Remote Inference] Supported params are ignored (#1562)

Author: kaisopos

src/oumi/core/inference/base_inference_engine.py

@@ -79,11 +79,20 @@ def infer(
                 "Only one of input or inference_config.input_path should be provided."
             )
 
-        if inference_config and inference_config.generation:
-            generation_params = inference_config.generation
-            self._check_unsupported_params(generation_params)
-        else:
-            generation_params = self._generation_params
+        # Ensure the inference config has up-to-date generation parameters.
+        if inference_config:
+            if inference_config.generation:
+                self._check_unsupported_params(inference_config.generation)
+            elif self._generation_params:
+                inference_config = copy.deepcopy(inference_config)
+                inference_config.generation = self._generation_params
+
+                # Warn the user: They provided an inference config without generation
+                # params, so what was the point of providing it in the first place?
+                logger.warning(
+                    "No generation parameters provided in the inference config. Using "
+                    "the generation parameters that the engine was initialized with."
+                )
 
         if input is not None:
             return self.infer_online(input, inference_config)

src/oumi/inference/anthropic_inference_engine.py

@@ -16,7 +16,7 @@
 
 from typing_extensions import override
 
-from oumi.core.configs import GenerationParams, RemoteParams
+from oumi.core.configs import GenerationParams, ModelParams, RemoteParams
 from oumi.core.types.conversation import Conversation, Message, Role
 from oumi.inference.remote_inference_engine import RemoteInferenceEngine
 from oumi.utils.logging import logger
@@ -55,7 +55,10 @@ def api_key_env_varname(self) -> Optional[str]:
 
     @override
     def _convert_conversation_to_api_input(
-        self, conversation: Conversation, generation_params: GenerationParams
+        self,
+        conversation: Conversation,
+        generation_params: GenerationParams,
+        model_params: ModelParams,
     ) -> dict[str, Any]:
         """Converts a conversation to an Anthropic API input.
 
@@ -68,6 +71,7 @@ def _convert_conversation_to_api_input(
         Args:
             conversation: The Oumi Conversation object to convert.
             generation_params: Parameters for text generation.
+            model_params: Model parameters to use during inference.
 
         Returns:
             Dict[str, Any]: A dictionary containing the formatted input for the
@@ -98,7 +102,7 @@ def _convert_conversation_to_api_input(
         # Build request body
         # See https://docs.anthropic.com/claude/reference/messages_post
         body = {
-            "model": self._model,
+            "model": model_params.model_name,
             "messages": self._get_list_of_message_json_dicts(
                 messages, group_adjacent_same_role_turns=True
             ),

src/oumi/inference/gcp_inference_engine.py

@@ -18,7 +18,7 @@
 import pydantic
 from typing_extensions import override
 
-from oumi.core.configs import GenerationParams, RemoteParams
+from oumi.core.configs import GenerationParams, ModelParams, RemoteParams
 from oumi.core.configs.params.guided_decoding_params import GuidedDecodingParams
 from oumi.core.types.conversation import Conversation
 from oumi.inference.remote_inference_engine import RemoteInferenceEngine
@@ -73,7 +73,10 @@ def _get_request_headers(
 
     @override
     def _convert_conversation_to_api_input(
-        self, conversation: Conversation, generation_params: GenerationParams
+        self,
+        conversation: Conversation,
+        generation_params: GenerationParams,
+        model_params: ModelParams,
     ) -> dict[str, Any]:
         """Converts a conversation to an OpenAI input.
 
@@ -82,12 +85,13 @@ def _convert_conversation_to_api_input(
         Args:
             conversation: The conversation to convert.
             generation_params: Parameters for generation during inference.
+            model_params: Model parameters to use during inference.
 
         Returns:
             Dict[str, Any]: A dictionary representing the Vertex input.
         """
         api_input = {
-            "model": self._model,
+            "model": model_params.model_name,
             "messages": self._get_list_of_message_json_dicts(
                 conversation.messages, group_adjacent_same_role_turns=True
             ),

src/oumi/inference/gemini_inference_engine.py

@@ -16,7 +16,7 @@
 
 from typing_extensions import override
 
-from oumi.core.configs import GenerationParams
+from oumi.core.configs import GenerationParams, ModelParams
 from oumi.core.types.conversation import Conversation
 from oumi.inference.gcp_inference_engine import (
     _convert_guided_decoding_config_to_api_input,
@@ -37,7 +37,10 @@ class GoogleGeminiInferenceEngine(RemoteInferenceEngine):
 
     @override
     def _convert_conversation_to_api_input(
-        self, conversation: Conversation, generation_params: GenerationParams
+        self,
+        conversation: Conversation,
+        generation_params: GenerationParams,
+        model_params: ModelParams,
     ) -> dict[str, Any]:
         """Converts a conversation to an Gemini API input.
 
@@ -46,12 +49,13 @@ def _convert_conversation_to_api_input(
         Args:
             conversation: The conversation to convert.
             generation_params: Parameters for generation during inference.
+            model_params: Model parameters to use during inference.
 
         Returns:
             Dict[str, Any]: A dictionary representing the Gemini input.
         """
         api_input = {
-            "model": self._model,
+            "model": model_params.model_name,
             "messages": self._get_list_of_message_json_dicts(
                 conversation.messages, group_adjacent_same_role_turns=True
             ),

src/oumi/inference/openai_inference_engine.py

@@ -12,10 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional
+import copy
+from typing import Any, Optional
 
 from typing_extensions import override
 
+from oumi.core.configs import GenerationParams, ModelParams
+from oumi.core.types.conversation import Conversation
 from oumi.inference.remote_inference_engine import RemoteInferenceEngine
 
 
@@ -33,3 +36,33 @@ def base_url(self) -> Optional[str]:
     def api_key_env_varname(self) -> Optional[str]:
         """Return the default environment variable name for the OpenAI API key."""
         return "OPENAI_API_KEY"
+
+    @override
+    def _convert_conversation_to_api_input(
+        self,
+        conversation: Conversation,
+        generation_params: GenerationParams,
+        model_params: ModelParams,
+    ) -> dict[str, Any]:
+        """Converts a conversation to an OpenAI input.
+
+        Documentation: https://platform.openai.com/docs/api-reference/chat/create
+
+        Args:
+            conversation: The conversation to convert.
+            generation_params: Parameters for generation during inference.
+            model_params: Model parameters to use during inference.
+
+        Returns:
+            Dict[str, Any]: A dictionary representing the OpenAI input.
+        """
+        # o1-preview does NOT support logit_bias.
+        if model_params.model_name == "o1-preview":
+            generation_params = copy.deepcopy(generation_params)
+            generation_params.logit_bias = {}
+
+        return super()._convert_conversation_to_api_input(
+            conversation=conversation,
+            generation_params=generation_params,
+            model_params=model_params,
+        )