diff --git a/google/genai/tests/tunings/test_tune.py b/google/genai/tests/tunings/test_tune.py index 6ef7fcbc8..4b0470270 100755 --- a/google/genai/tests/tunings/test_tune.py +++ b/google/genai/tests/tunings/test_tune.py @@ -65,6 +65,37 @@ ), exception_if_mldev="is not supported in Gemini API", ), + pytest_helper.TestTableItem( + name="test_tune_simple_dpo", + parameters=genai_types.CreateTuningJobParameters( + base_model="gemini-2.5-flash", + training_dataset=genai_types.TuningDataset( + gcs_uri="gs://cloud-samples-data/ai-platform/generative_ai/gemini-1_5/text/sft_train_data.jsonl", + ), + config=genai_types.CreateTuningJobConfig( + tuned_model_display_name="Model display name", + epoch_count=1, + method="PREFERENCE_TUNING", + ), + ), + exception_if_mldev="parameter is not supported in Gemini API.", + ), + pytest_helper.TestTableItem( + name="test_tune_dpo_with_beta", + parameters=genai_types.CreateTuningJobParameters( + base_model="gemini-2.5-flash", + training_dataset=genai_types.TuningDataset( + gcs_uri="gs://cloud-samples-data/ai-platform/generative_ai/gemini-1_5/text/sft_train_data.jsonl", + ), + config=genai_types.CreateTuningJobConfig( + tuned_model_display_name="Model display name", + epoch_count=1, + method=genai_types.TuningMethod.PREFERENCE_TUNING, + beta=0.5, + ), + ), + exception_if_mldev="parameter is not supported in Gemini API.", + ), pytest_helper.TestTableItem( name="test_non_pretuned_model_with_checkpoint_id", parameters=genai_types.CreateTuningJobParameters( diff --git a/google/genai/tunings.py b/google/genai/tunings.py index 784556a35..119e04362 100644 --- a/google/genai/tunings.py +++ b/google/genai/tunings.py @@ -128,6 +128,9 @@ def _CreateTuningJobConfig_to_mldev( if getv(from_object, ['labels']) is not None: raise ValueError('labels parameter is not supported in Gemini API.') + if getv(from_object, ['beta']) is not None: + raise ValueError('beta parameter is not supported in Gemini API.') + return to_object @@ -138,14 +141,28 @@ def _CreateTuningJobConfig_to_vertex( ) -> dict[str, Any]: to_object: dict[str, Any] = {} - if getv(from_object, ['validation_dataset']) is not None: - setv( - parent_object, - ['supervisedTuningSpec'], - _TuningValidationDataset_to_vertex( - getv(from_object, ['validation_dataset']), to_object, root_object - ), - ) + discriminator = getv(root_object, ['config', 'method']) + if discriminator is None: + discriminator = 'SUPERVISED_FINE_TUNING' + if discriminator == 'SUPERVISED_FINE_TUNING': + if getv(from_object, ['validation_dataset']) is not None: + setv( + parent_object, + ['supervisedTuningSpec'], + _TuningValidationDataset_to_vertex( + getv(from_object, ['validation_dataset']), to_object, root_object + ), + ) + + elif discriminator == 'PREFERENCE_TUNING': + if getv(from_object, ['validation_dataset']) is not None: + setv( + parent_object, + ['preferenceOptimizationSpec'], + _TuningValidationDataset_to_vertex( + getv(from_object, ['validation_dataset']), to_object, root_object + ), + ) if getv(from_object, ['tuned_model_display_name']) is not None: setv( @@ -157,33 +174,85 @@ def _CreateTuningJobConfig_to_vertex( if getv(from_object, ['description']) is not None: setv(parent_object, ['description'], getv(from_object, ['description'])) - if getv(from_object, ['epoch_count']) is not None: - setv( - parent_object, - ['supervisedTuningSpec', 'hyperParameters', 'epochCount'], - getv(from_object, ['epoch_count']), - ) + discriminator = getv(root_object, ['config', 'method']) + if discriminator is None: + discriminator = 'SUPERVISED_FINE_TUNING' + if discriminator == 'SUPERVISED_FINE_TUNING': + if getv(from_object, ['epoch_count']) is not None: + setv( + parent_object, + ['supervisedTuningSpec', 'hyperParameters', 'epochCount'], + getv(from_object, ['epoch_count']), + ) - if getv(from_object, ['learning_rate_multiplier']) is not None: - setv( - parent_object, - ['supervisedTuningSpec', 'hyperParameters', 'learningRateMultiplier'], - getv(from_object, ['learning_rate_multiplier']), - ) + elif discriminator == 'PREFERENCE_TUNING': + if getv(from_object, ['epoch_count']) is not None: + setv( + parent_object, + ['preferenceOptimizationSpec', 'hyperParameters', 'epochCount'], + getv(from_object, ['epoch_count']), + ) - if getv(from_object, ['export_last_checkpoint_only']) is not None: - setv( - parent_object, - ['supervisedTuningSpec', 'exportLastCheckpointOnly'], - getv(from_object, ['export_last_checkpoint_only']), - ) + discriminator = getv(root_object, ['config', 'method']) + if discriminator is None: + discriminator = 'SUPERVISED_FINE_TUNING' + if discriminator == 'SUPERVISED_FINE_TUNING': + if getv(from_object, ['learning_rate_multiplier']) is not None: + setv( + parent_object, + ['supervisedTuningSpec', 'hyperParameters', 'learningRateMultiplier'], + getv(from_object, ['learning_rate_multiplier']), + ) - if getv(from_object, ['adapter_size']) is not None: - setv( - parent_object, - ['supervisedTuningSpec', 'hyperParameters', 'adapterSize'], - getv(from_object, ['adapter_size']), - ) + elif discriminator == 'PREFERENCE_TUNING': + if getv(from_object, ['learning_rate_multiplier']) is not None: + setv( + parent_object, + [ + 'preferenceOptimizationSpec', + 'hyperParameters', + 'learningRateMultiplier', + ], + getv(from_object, ['learning_rate_multiplier']), + ) + + discriminator = getv(root_object, ['config', 'method']) + if discriminator is None: + discriminator = 'SUPERVISED_FINE_TUNING' + if discriminator == 'SUPERVISED_FINE_TUNING': + if getv(from_object, ['export_last_checkpoint_only']) is not None: + setv( + parent_object, + ['supervisedTuningSpec', 'exportLastCheckpointOnly'], + getv(from_object, ['export_last_checkpoint_only']), + ) + + elif discriminator == 'PREFERENCE_TUNING': + if getv(from_object, ['export_last_checkpoint_only']) is not None: + setv( + parent_object, + ['preferenceOptimizationSpec', 'exportLastCheckpointOnly'], + getv(from_object, ['export_last_checkpoint_only']), + ) + + discriminator = getv(root_object, ['config', 'method']) + if discriminator is None: + discriminator = 'SUPERVISED_FINE_TUNING' + if discriminator == 'SUPERVISED_FINE_TUNING': + if getv(from_object, ['adapter_size']) is not None: + setv( + parent_object, + ['supervisedTuningSpec', 'hyperParameters', 'adapterSize'], + getv(from_object, ['adapter_size']), + ) + + elif discriminator == 'PREFERENCE_TUNING': + if getv(from_object, ['adapter_size']) is not None: + setv( + parent_object, + ['preferenceOptimizationSpec', 'hyperParameters', 'adapterSize'], + getv(from_object, ['adapter_size']), + ) if getv(from_object, ['batch_size']) is not None: raise ValueError('batch_size parameter is not supported in Vertex AI.') @@ -191,18 +260,39 @@ def _CreateTuningJobConfig_to_vertex( if getv(from_object, ['learning_rate']) is not None: raise ValueError('learning_rate parameter is not supported in Vertex AI.') - if getv(from_object, ['evaluation_config']) is not None: - setv( - parent_object, - ['supervisedTuningSpec', 'evaluationConfig'], - _EvaluationConfig_to_vertex( - getv(from_object, ['evaluation_config']), to_object, root_object - ), - ) + discriminator = getv(root_object, ['config', 'method']) + if discriminator is None: + discriminator = 'SUPERVISED_FINE_TUNING' + if discriminator == 'SUPERVISED_FINE_TUNING': + if getv(from_object, ['evaluation_config']) is not None: + setv( + parent_object, + ['supervisedTuningSpec', 'evaluationConfig'], + _EvaluationConfig_to_vertex( + getv(from_object, ['evaluation_config']), to_object, root_object + ), + ) + + elif discriminator == 'PREFERENCE_TUNING': + if getv(from_object, ['evaluation_config']) is not None: + setv( + parent_object, + ['preferenceOptimizationSpec', 'evaluationConfig'], + _EvaluationConfig_to_vertex( + getv(from_object, ['evaluation_config']), to_object, root_object + ), + ) if getv(from_object, ['labels']) is not None: setv(parent_object, ['labels'], getv(from_object, ['labels'])) + if getv(from_object, ['beta']) is not None: + setv( + parent_object, + ['preferenceOptimizationSpec', 'hyperParameters', 'beta'], + getv(from_object, ['beta']), + ) + return to_object @@ -219,12 +309,8 @@ def _CreateTuningJobParametersPrivate_to_mldev( setv(to_object, ['preTunedModel'], getv(from_object, ['pre_tuned_model'])) if getv(from_object, ['training_dataset']) is not None: - setv( - to_object, - ['tuningTask', 'trainingData'], - _TuningDataset_to_mldev( - getv(from_object, ['training_dataset']), to_object, root_object - ), + _TuningDataset_to_mldev( + getv(from_object, ['training_dataset']), to_object, root_object ) if getv(from_object, ['config']) is not None: @@ -501,19 +587,44 @@ def _TuningDataset_to_vertex( root_object: Optional[Union[dict[str, Any], object]] = None, ) -> dict[str, Any]: to_object: dict[str, Any] = {} - if getv(from_object, ['gcs_uri']) is not None: - setv( - parent_object, - ['supervisedTuningSpec', 'trainingDatasetUri'], - getv(from_object, ['gcs_uri']), - ) - if getv(from_object, ['vertex_dataset_resource']) is not None: - setv( - parent_object, - ['supervisedTuningSpec', 'trainingDatasetUri'], - getv(from_object, ['vertex_dataset_resource']), - ) + discriminator = getv(root_object, ['config', 'method']) + if discriminator is None: + discriminator = 'SUPERVISED_FINE_TUNING' + if discriminator == 'SUPERVISED_FINE_TUNING': + if getv(from_object, ['gcs_uri']) is not None: + setv( + parent_object, + ['supervisedTuningSpec', 'trainingDatasetUri'], + getv(from_object, ['gcs_uri']), + ) + + elif discriminator == 'PREFERENCE_TUNING': + if getv(from_object, ['gcs_uri']) is not None: + setv( + parent_object, + ['preferenceOptimizationSpec', 'trainingDatasetUri'], + getv(from_object, ['gcs_uri']), + ) + + discriminator = getv(root_object, ['config', 'method']) + if discriminator is None: + discriminator = 'SUPERVISED_FINE_TUNING' + if discriminator == 'SUPERVISED_FINE_TUNING': + if getv(from_object, ['vertex_dataset_resource']) is not None: + setv( + parent_object, + ['supervisedTuningSpec', 'trainingDatasetUri'], + getv(from_object, ['vertex_dataset_resource']), + ) + + elif discriminator == 'PREFERENCE_TUNING': + if getv(from_object, ['vertex_dataset_resource']) is not None: + setv( + parent_object, + ['preferenceOptimizationSpec', 'trainingDatasetUri'], + getv(from_object, ['vertex_dataset_resource']), + ) if getv(from_object, ['examples']) is not None: raise ValueError('examples parameter is not supported in Vertex AI.') @@ -635,6 +746,13 @@ def _TuningJob_from_vertex( getv(from_object, ['supervisedTuningSpec']), ) + if getv(from_object, ['preferenceOptimizationSpec']) is not None: + setv( + to_object, + ['preference_optimization_spec'], + getv(from_object, ['preferenceOptimizationSpec']), + ) + if getv(from_object, ['tuningDataStats']) is not None: setv( to_object, ['tuning_data_stats'], getv(from_object, ['tuningDataStats']) @@ -950,7 +1068,7 @@ def _tune( training_dataset: types.TuningDatasetOrDict, config: Optional[types.CreateTuningJobConfigOrDict] = None, ) -> types.TuningJob: - """Creates a supervised fine-tuning job and returns the TuningJob object. + """Creates a tuning job and returns the TuningJob object. Args: base_model: The name of the model to tune. @@ -1023,7 +1141,7 @@ def _tune_mldev( training_dataset: types.TuningDatasetOrDict, config: Optional[types.CreateTuningJobConfigOrDict] = None, ) -> types.TuningOperation: - """Creates a supervised fine-tuning job and returns the TuningJob object. + """Creates a tuning job and returns the TuningJob object. Args: base_model: The name of the model to tune. @@ -1419,7 +1537,7 @@ async def _tune( training_dataset: types.TuningDatasetOrDict, config: Optional[types.CreateTuningJobConfigOrDict] = None, ) -> types.TuningJob: - """Creates a supervised fine-tuning job and returns the TuningJob object. + """Creates a tuning job and returns the TuningJob object. Args: base_model: The name of the model to tune. @@ -1492,7 +1610,7 @@ async def _tune_mldev( training_dataset: types.TuningDatasetOrDict, config: Optional[types.CreateTuningJobConfigOrDict] = None, ) -> types.TuningOperation: - """Creates a supervised fine-tuning job and returns the TuningJob object. + """Creates a tuning job and returns the TuningJob object. Args: base_model: The name of the model to tune. diff --git a/google/genai/types.py b/google/genai/types.py index 1318e5ac8..b121b541e 100644 --- a/google/genai/types.py +++ b/google/genai/types.py @@ -704,6 +704,15 @@ class VideoCompressionQuality(_common.CaseInSensitiveEnum): with a larger file size.""" +class TuningMethod(_common.CaseInSensitiveEnum): + """Enum representing the tuning method.""" + + SUPERVISED_FINE_TUNING = 'SUPERVISED_FINE_TUNING' + """Supervised fine tuning.""" + PREFERENCE_TUNING = 'PREFERENCE_TUNING' + """Preference optimization tuning.""" + + class FileState(_common.CaseInSensitiveEnum): """State for the lifecycle of a File.""" @@ -9659,6 +9668,90 @@ class SupervisedTuningSpecDict(TypedDict, total=False): ] +class PreferenceOptimizationHyperParameters(_common.BaseModel): + """Hyperparameters for Preference Optimization. + + This data type is not supported in Gemini API. + """ + + adapter_size: Optional[AdapterSize] = Field( + default=None, + description="""Optional. Adapter size for preference optimization.""", + ) + beta: Optional[float] = Field( + default=None, + description="""Optional. Weight for KL Divergence regularization.""", + ) + epoch_count: Optional[int] = Field( + default=None, + description="""Optional. Number of complete passes the model makes over the entire training dataset during training.""", + ) + learning_rate_multiplier: Optional[float] = Field( + default=None, + description="""Optional. Multiplier for adjusting the default learning rate.""", + ) + + +class PreferenceOptimizationHyperParametersDict(TypedDict, total=False): + """Hyperparameters for Preference Optimization. + + This data type is not supported in Gemini API. + """ + + adapter_size: Optional[AdapterSize] + """Optional. Adapter size for preference optimization.""" + + beta: Optional[float] + """Optional. Weight for KL Divergence regularization.""" + + epoch_count: Optional[int] + """Optional. Number of complete passes the model makes over the entire training dataset during training.""" + + learning_rate_multiplier: Optional[float] + """Optional. Multiplier for adjusting the default learning rate.""" + + +PreferenceOptimizationHyperParametersOrDict = Union[ + PreferenceOptimizationHyperParameters, + PreferenceOptimizationHyperParametersDict, +] + + +class PreferenceOptimizationSpec(_common.BaseModel): + """Preference optimization tuning spec for tuning.""" + + hyper_parameters: Optional[PreferenceOptimizationHyperParameters] = Field( + default=None, + description="""Optional. Hyperparameters for Preference Optimization.""", + ) + training_dataset_uri: Optional[str] = Field( + default=None, + description="""Required. Cloud Storage path to file containing training dataset for preference optimization tuning. The dataset must be formatted as a JSONL file.""", + ) + validation_dataset_uri: Optional[str] = Field( + default=None, + description="""Optional. Cloud Storage path to file containing validation dataset for preference optimization tuning. The dataset must be formatted as a JSONL file.""", + ) + + +class PreferenceOptimizationSpecDict(TypedDict, total=False): + """Preference optimization tuning spec for tuning.""" + + hyper_parameters: Optional[PreferenceOptimizationHyperParametersDict] + """Optional. Hyperparameters for Preference Optimization.""" + + training_dataset_uri: Optional[str] + """Required. Cloud Storage path to file containing training dataset for preference optimization tuning. The dataset must be formatted as a JSONL file.""" + + validation_dataset_uri: Optional[str] + """Optional. Cloud Storage path to file containing validation dataset for preference optimization tuning. The dataset must be formatted as a JSONL file.""" + + +PreferenceOptimizationSpecOrDict = Union[ + PreferenceOptimizationSpec, PreferenceOptimizationSpecDict +] + + class GcsDestination(_common.BaseModel): """The Google Cloud Storage location where the output is to be written to.""" @@ -10844,6 +10937,9 @@ class TuningJob(_common.BaseModel): supervised_tuning_spec: Optional[SupervisedTuningSpec] = Field( default=None, description="""Tuning Spec for Supervised Fine Tuning.""" ) + preference_optimization_spec: Optional[PreferenceOptimizationSpec] = Field( + default=None, description="""Tuning Spec for Preference Optimization.""" + ) tuning_data_stats: Optional[TuningDataStats] = Field( default=None, description="""Output only. The tuning data statistics associated with this TuningJob.""", @@ -10944,6 +11040,9 @@ class TuningJobDict(TypedDict, total=False): supervised_tuning_spec: Optional[SupervisedTuningSpecDict] """Tuning Spec for Supervised Fine Tuning.""" + preference_optimization_spec: Optional[PreferenceOptimizationSpecDict] + """Tuning Spec for Preference Optimization.""" + tuning_data_stats: Optional[TuningDataStatsDict] """Output only. The tuning data statistics associated with this TuningJob.""" @@ -11205,11 +11304,15 @@ class TuningValidationDatasetDict(TypedDict, total=False): class CreateTuningJobConfig(_common.BaseModel): - """Supervised fine-tuning job creation request - optional fields.""" + """Fine-tuning job creation request - optional fields.""" http_options: Optional[HttpOptions] = Field( default=None, description="""Used to override HTTP request options.""" ) + method: Optional[TuningMethod] = Field( + default=None, + description="""The method to use for tuning (SUPERVISED_FINE_TUNING or PREFERENCE_TUNING). If not set, the default method (SFT) will be used.""", + ) validation_dataset: Optional[TuningValidationDataset] = Field( default=None, description="""Validation dataset for tuning. The dataset must be formatted as a JSONL file.""", @@ -11231,7 +11334,7 @@ class CreateTuningJobConfig(_common.BaseModel): ) export_last_checkpoint_only: Optional[bool] = Field( default=None, - description="""If set to true, disable intermediate checkpoints for SFT and only the last checkpoint will be exported. Otherwise, enable intermediate checkpoints for SFT.""", + description="""If set to true, disable intermediate checkpoints and only the last checkpoint will be exported. Otherwise, enable intermediate checkpoints.""", ) pre_tuned_model_checkpoint_id: Optional[str] = Field( default=None, @@ -11255,14 +11358,21 @@ class CreateTuningJobConfig(_common.BaseModel): default=None, description="""Optional. The labels with user-defined metadata to organize TuningJob and generated resources such as Model and Endpoint. Label keys and values can be no longer than 64 characters (Unicode codepoints), can only contain lowercase letters, numeric characters, underscores and dashes. International characters are allowed. See https://goo.gl/xmQnxf for more information and examples of labels.""", ) + beta: Optional[float] = Field( + default=None, + description="""Weight for KL Divergence regularization, Preference Optimization tuning only.""", + ) class CreateTuningJobConfigDict(TypedDict, total=False): - """Supervised fine-tuning job creation request - optional fields.""" + """Fine-tuning job creation request - optional fields.""" http_options: Optional[HttpOptionsDict] """Used to override HTTP request options.""" + method: Optional[TuningMethod] + """The method to use for tuning (SUPERVISED_FINE_TUNING or PREFERENCE_TUNING). If not set, the default method (SFT) will be used.""" + validation_dataset: Optional[TuningValidationDatasetDict] """Validation dataset for tuning. The dataset must be formatted as a JSONL file.""" @@ -11279,7 +11389,7 @@ class CreateTuningJobConfigDict(TypedDict, total=False): """Multiplier for adjusting the default learning rate.""" export_last_checkpoint_only: Optional[bool] - """If set to true, disable intermediate checkpoints for SFT and only the last checkpoint will be exported. Otherwise, enable intermediate checkpoints for SFT.""" + """If set to true, disable intermediate checkpoints and only the last checkpoint will be exported. Otherwise, enable intermediate checkpoints.""" pre_tuned_model_checkpoint_id: Optional[str] """The optional checkpoint id of the pre-tuned model to use for tuning, if applicable.""" @@ -11299,6 +11409,9 @@ class CreateTuningJobConfigDict(TypedDict, total=False): labels: Optional[dict[str, str]] """Optional. The labels with user-defined metadata to organize TuningJob and generated resources such as Model and Endpoint. Label keys and values can be no longer than 64 characters (Unicode codepoints), can only contain lowercase letters, numeric characters, underscores and dashes. International characters are allowed. See https://goo.gl/xmQnxf for more information and examples of labels.""" + beta: Optional[float] + """Weight for KL Divergence regularization, Preference Optimization tuning only.""" + CreateTuningJobConfigOrDict = Union[ CreateTuningJobConfig, CreateTuningJobConfigDict @@ -11306,7 +11419,7 @@ class CreateTuningJobConfigDict(TypedDict, total=False): class _CreateTuningJobParametersPrivate(_common.BaseModel): - """Supervised fine-tuning job creation parameters - optional fields.""" + """Fine-tuning job creation parameters - optional fields.""" base_model: Optional[str] = Field( default=None, @@ -11325,7 +11438,7 @@ class _CreateTuningJobParametersPrivate(_common.BaseModel): class _CreateTuningJobParametersPrivateDict(TypedDict, total=False): - """Supervised fine-tuning job creation parameters - optional fields.""" + """Fine-tuning job creation parameters - optional fields.""" base_model: Optional[str] """The base model that is being tuned, e.g., "gemini-2.5-flash".""" @@ -16148,7 +16261,7 @@ class ComputeTokensResultDict(TypedDict, total=False): class CreateTuningJobParameters(_common.BaseModel): - """Supervised fine-tuning job creation parameters - optional fields.""" + """Fine-tuning job creation parameters - optional fields.""" base_model: Optional[str] = Field( default=None, @@ -16164,7 +16277,7 @@ class CreateTuningJobParameters(_common.BaseModel): class CreateTuningJobParametersDict(TypedDict, total=False): - """Supervised fine-tuning job creation parameters - optional fields.""" + """Fine-tuning job creation parameters - optional fields.""" base_model: Optional[str] """The base model that is being tuned, e.g., "gemini-2.5-flash"."""