Skip to content

Commit a433de2

Browse files
cleop-googlecopybara-github
cleop-google
authored and
copybara-github
committed
chore: migrate to the read config field in multimodal dataset API requests
PiperOrigin-RevId: 758602332
1 parent c551cc6 commit a433de2

2 files changed

Lines changed: 63 additions & 43 deletions

File tree

‎google/cloud/aiplatform/preview/datasets.py‎

Lines changed: 45 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -1371,14 +1371,12 @@ def assemble(
13711371
load_dataframe is True, otherwise None.
13721372
"""
13731373
bigframes = _try_import_bigframes()
1374-
request = gca_dataset_service.AssembleDataRequest(name=self.resource_name)
1375-
if self.request_column_name is not None:
1376-
request.request_column_name = self.request_column_name
1377-
else:
1378-
template_config_to_use = _resolve_template_config(self, template_config)
1379-
request.gemini_template_config = (
1380-
template_config_to_use._raw_gemini_template_config
1381-
)
1374+
request = gca_dataset_service.AssembleDataRequest(
1375+
name=self.resource_name,
1376+
gemini_request_read_config=self._build_gemini_request_read_config(
1377+
template_config
1378+
),
1379+
)
13821380

13831381
assemble_lro = self.api_client.assemble_data(
13841382
request=request, timeout=assemble_request_timeout
@@ -1422,7 +1420,7 @@ def assess_tuning_resources(
14221420
dataset.
14231421
14241422
"""
1425-
request = _build_assess_data_request(self, template_config)
1423+
request = self._build_assess_data_request(template_config)
14261424
request.tuning_resource_usage_assessment_config = (
14271425
gca_dataset_service.AssessDataRequest.TuningResourceUsageAssessmentConfig(
14281426
model_name=model_name
@@ -1486,7 +1484,7 @@ def assess_tuning_validity(
14861484
if dataset_usage_enum == DatasetUsage.DATASET_USAGE_UNSPECIFIED:
14871485
raise ValueError("Dataset usage must be specified.")
14881486

1489-
request = _build_assess_data_request(self, template_config)
1487+
request = self._build_assess_data_request(template_config)
14901488
request.tuning_validation_assessment_config = (
14911489
gca_dataset_service.AssessDataRequest.TuningValidationAssessmentConfig(
14921490
model_name=model_name,
@@ -1501,32 +1499,42 @@ def assess_tuning_validity(
15011499
errors=assessment_result.tuning_validation_assessment_result.errors
15021500
)
15031501

1502+
def _build_assess_data_request(
1503+
self,
1504+
template_config: Optional[GeminiTemplateConfig] = None,
1505+
):
1506+
return gca_dataset_service.AssessDataRequest(
1507+
name=self.resource_name,
1508+
gemini_request_read_config=self._build_gemini_request_read_config(
1509+
template_config
1510+
),
1511+
)
15041512

1505-
def _resolve_template_config(
1506-
dataset: MultimodalDataset,
1507-
template_config: Optional[GeminiTemplateConfig] = None,
1508-
) -> GeminiTemplateConfig:
1509-
"""Returns the passed template config if it is not None, otherwise
1510-
returns the template config attached to the dataset.
1511-
"""
1512-
if template_config is not None:
1513-
return template_config
1514-
elif dataset.template_config is not None:
1515-
return dataset.template_config
1516-
else:
1517-
raise ValueError("No template config was passed or attached to the dataset.")
1518-
1513+
def _build_gemini_request_read_config(
1514+
self, template_config: Optional[GeminiTemplateConfig] = None
1515+
):
1516+
if self.request_column_name is not None:
1517+
return gca_dataset_service.GeminiRequestReadConfig(
1518+
assembled_request_column_name=self.request_column_name
1519+
)
1520+
else:
1521+
template_config_to_use = self._resolve_template_config(template_config)
1522+
return gca_dataset_service.GeminiRequestReadConfig(
1523+
template_config=template_config_to_use._raw_gemini_template_config
1524+
)
15191525

1520-
def _build_assess_data_request(
1521-
dataset: MultimodalDataset,
1522-
template_config: Optional[GeminiTemplateConfig] = None,
1523-
):
1524-
request = gca_dataset_service.AssessDataRequest(name=dataset.resource_name)
1525-
if dataset.request_column_name is not None:
1526-
request.request_column_name = dataset.request_column_name
1527-
else:
1528-
template_config_to_use = _resolve_template_config(dataset, template_config)
1529-
request.gemini_template_config = (
1530-
template_config_to_use._raw_gemini_template_config
1531-
)
1532-
return request
1526+
def _resolve_template_config(
1527+
self,
1528+
template_config: Optional[GeminiTemplateConfig] = None,
1529+
) -> GeminiTemplateConfig:
1530+
"""Returns the passed template config if it is not None, otherwise
1531+
returns the template config attached to the dataset.
1532+
"""
1533+
if template_config is not None:
1534+
return template_config
1535+
elif self.template_config is not None:
1536+
return self.template_config
1537+
else:
1538+
raise ValueError(
1539+
"No template config was passed or attached to the dataset."
1540+
)

‎tests/unit/aiplatform/test_multimodal_datasets.py‎

Lines changed: 18 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -675,7 +675,9 @@ def test_assess_tuning_resources(self, assess_data_tuning_resources_mock):
675675
tuning_resource_usage_assessment_config=gca_dataset_service.AssessDataRequest.TuningResourceUsageAssessmentConfig(
676676
model_name="gemini-1.5-flash-exp"
677677
),
678-
gemini_template_config=template_config._raw_gemini_template_config,
678+
gemini_request_read_config=gca_dataset_service.GeminiRequestReadConfig(
679+
template_config=template_config._raw_gemini_template_config
680+
),
679681
),
680682
timeout=None,
681683
)
@@ -696,7 +698,9 @@ def test_assess_tuning_resources_request_column_name(
696698
tuning_resource_usage_assessment_config=gca_dataset_service.AssessDataRequest.TuningResourceUsageAssessmentConfig(
697699
model_name="gemini-1.5-flash-exp"
698700
),
699-
request_column_name="requests",
701+
gemini_request_read_config=gca_dataset_service.GeminiRequestReadConfig(
702+
assembled_request_column_name="requests"
703+
),
700704
),
701705
timeout=None,
702706
)
@@ -720,7 +724,9 @@ def test_assess_tuning_validity(self, assess_data_tuning_validation_mock):
720724
model_name="gemini-1.5-flash-exp",
721725
dataset_usage=gca_dataset_service.AssessDataRequest.TuningValidationAssessmentConfig.DatasetUsage.SFT_TRAINING,
722726
),
723-
gemini_template_config=template_config._raw_gemini_template_config,
727+
gemini_request_read_config=gca_dataset_service.GeminiRequestReadConfig(
728+
template_config=template_config._raw_gemini_template_config
729+
),
724730
),
725731
timeout=None,
726732
)
@@ -743,7 +749,9 @@ def test_assess_tuning_validity_request_column_name(
743749
model_name="gemini-1.5-flash-exp",
744750
dataset_usage=gca_dataset_service.AssessDataRequest.TuningValidationAssessmentConfig.DatasetUsage.SFT_TRAINING,
745751
),
746-
request_column_name="requests",
752+
gemini_request_read_config=gca_dataset_service.GeminiRequestReadConfig(
753+
assembled_request_column_name="requests"
754+
),
747755
),
748756
timeout=None,
749757
)
@@ -783,7 +791,9 @@ def test_assemble(self, assemble_data_mock):
783791
assemble_data_mock.assert_called_once_with(
784792
request=gca_dataset_service.AssembleDataRequest(
785793
name=_TEST_NAME,
786-
gemini_template_config=template_config._raw_gemini_template_config,
794+
gemini_request_read_config=gca_dataset_service.GeminiRequestReadConfig(
795+
template_config=template_config._raw_gemini_template_config
796+
),
787797
),
788798
timeout=None,
789799
)
@@ -802,7 +812,9 @@ def test_assemble_request_column_name(self, assemble_data_mock):
802812
assemble_data_mock.assert_called_once_with(
803813
request=gca_dataset_service.AssembleDataRequest(
804814
name=_TEST_NAME,
805-
request_column_name="requests",
815+
gemini_request_read_config=gca_dataset_service.GeminiRequestReadConfig(
816+
assembled_request_column_name="requests"
817+
),
806818
),
807819
timeout=None,
808820
)

0 commit comments

Comments
 (0)