diff --git a/litellm/llms/fine_tuning_apis/vertex_ai.py b/litellm/llms/fine_tuning_apis/vertex_ai.py index 0fd882c815e..4c5938e68e9 100644 --- a/litellm/llms/fine_tuning_apis/vertex_ai.py +++ b/litellm/llms/fine_tuning_apis/vertex_ai.py @@ -88,19 +88,28 @@ class VertexFineTuningAPI(VertexLLM): def convert_openai_request_to_vertex( self, create_fine_tuning_job_data: FineTuningJobCreate, **kwargs - ) -> dict: + ) -> FineTuningJobCreate: """ convert request from OpenAI format to Vertex format https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/tuning supervised_tuning_spec = FineTunesupervisedTuningSpec( """ + hyperparameters = create_fine_tuning_job_data.hyperparameters supervised_tuning_spec = FineTunesupervisedTuningSpec( training_dataset_uri=create_fine_tuning_job_data.training_file, validation_dataset=create_fine_tuning_job_data.validation_file, - epoch_count=create_fine_tuning_job_data.hyperparameters.n_epochs, - learning_rate_multiplier=create_fine_tuning_job_data.hyperparameters.learning_rate_multiplier, - adapter_size=kwargs.get("AdapterSize"), ) + + if hyperparameters: + if hyperparameters.n_epochs: + supervised_tuning_spec["epoch_count"] = int(hyperparameters.n_epochs) + if hyperparameters.learning_rate_multiplier: + supervised_tuning_spec["learning_rate_multiplier"] = float( + hyperparameters.learning_rate_multiplier + ) + + supervised_tuning_spec["adapter_size"] = kwargs.get("adapter_size") + fine_tune_job = FineTuneJobCreate( baseModel=create_fine_tuning_job_data.model, supervisedTuningSpec=supervised_tuning_spec, @@ -113,7 +122,7 @@ class VertexFineTuningAPI(VertexLLM): self, fine_tuning_url: str, headers: dict, - request_data: FineTuneJobCreate, + request_data: FineTuningJobCreate, ): from litellm.fine_tuning.main import FineTuningJob diff --git a/litellm/tests/test_fine_tuning_api.py b/litellm/tests/test_fine_tuning_api.py index 6adb5e36906..4eebcd2f7fd 100644 --- a/litellm/tests/test_fine_tuning_api.py +++ b/litellm/tests/test_fine_tuning_api.py @@ -256,7 +256,7 @@ def test_convert_openai_request_to_vertex_with_adapter_size(): ) result = vertex_finetune_api.convert_openai_request_to_vertex( - openai_data, AdapterSize="SMALL" + openai_data, adapter_size="SMALL" ) print("converted vertex ai result=", result)