Skip to content

Commit 87a86af

Browse files
BenjaminKazemicopybara-github
BenjaminKazemi
authored and
copybara-github
committed
feat: GenAI SDK client(multimodal) - Support Assess Batch Prediction Resources for the multimodal datasets.
PiperOrigin-RevId: 868355650
1 parent 6907f89 commit 87a86af

File tree

2 files changed

+173
-0
lines changed

2 files changed

+173
-0
lines changed

‎tests/unit/vertexai/genai/replays/test_assess_multimodal_dataset.py‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,28 @@ def test_assess_tuning_resources(client):
6666
assert isinstance(response, types.TuningResourceUsageAssessmentResult)
6767

6868

69+
def test_assess_batch_prediction_resources(client):
70+
response = client.datasets.assess_batch_prediction_resources(
71+
dataset_name=DATASET,
72+
model_name="gemini-2.5-flash-001",
73+
template_config=types.GeminiTemplateConfig(
74+
gemini_example=types.GeminiExample(
75+
contents=[
76+
{
77+
"role": "user",
78+
"parts": [{"text": "What is the capital of {name}?"}],
79+
},
80+
{
81+
"role": "model",
82+
"parts": [{"text": "{capital}"}],
83+
},
84+
],
85+
),
86+
),
87+
)
88+
assert isinstance(response, types.BatchPredictionResourceUsageAssessmentResult)
89+
90+
6991
pytestmark = pytest_helper.setup(
7092
file=__file__,
7193
globals_for_file=globals(),
@@ -114,3 +136,26 @@ async def test_assess_tuning_resources_async(client):
114136
),
115137
)
116138
assert isinstance(response, types.TuningResourceUsageAssessmentResult)
139+
140+
141+
@pytest.mark.asyncio
142+
async def test_assess_batch_prediction_resources_async(client):
143+
response = await client.aio.datasets.assess_batch_prediction_resources(
144+
dataset_name=DATASET,
145+
model_name="gemini-2.5-flash-001",
146+
template_config=types.GeminiTemplateConfig(
147+
gemini_example=types.GeminiExample(
148+
contents=[
149+
{
150+
"role": "user",
151+
"parts": [{"text": "What is the capital of {name}?"}],
152+
},
153+
{
154+
"role": "model",
155+
"parts": [{"text": "{capital}"}],
156+
},
157+
],
158+
),
159+
),
160+
)
161+
assert isinstance(response, types.BatchPredictionResourceUsageAssessmentResult)

‎vertexai/_genai/datasets.py‎

Lines changed: 128 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1054,6 +1054,70 @@ def assess_tuning_resources(
10541054
response["tuningResourceUsageAssessmentResult"],
10551055
)
10561056

1057+
def assess_batch_prediction_resources(
1058+
self,
1059+
*,
1060+
dataset_name: str,
1061+
model_name: str,
1062+
template_config: Optional[types.GeminiTemplateConfigOrDict] = None,
1063+
config: Optional[types.AssessDatasetConfigOrDict] = None,
1064+
) -> types.BatchPredictionResourceUsageAssessmentResult:
1065+
"""Assess the batch prediction resources required for a given model.
1066+
1067+
Args:
1068+
dataset_name:
1069+
Required. The name of the dataset to assess the batch prediction
1070+
resources.
1071+
model_name:
1072+
Required. The name of the model to assess the batch prediction
1073+
resources.
1074+
template_config:
1075+
Optional. The template config used to assemble the dataset
1076+
before assessing the batch prediction resources. If not provided,
1077+
the template config attached to the dataset will be used. Required
1078+
if no template config is attached to the dataset.
1079+
template_config:
1080+
Optional. The template config used to assemble the dataset
1081+
before assessing the batch prediction resources. If not provided, the
1082+
template config attached to the dataset will be used. Required
1083+
if no template config is attached to the dataset.
1084+
config:
1085+
Optional. A configuration for assessing the batch prediction
1086+
resources. If not provided, the default configuration will be
1087+
used.
1088+
1089+
Returns:
1090+
A types.BatchPredictionResourceUsageAssessmentResult object
1091+
representing the batch prediction resource usage assessment result.
1092+
It contains the following keys:
1093+
- token_count: The number of tokens in the dataset.
1094+
- audio_token_count: The number of audio tokens in the dataset.
1095+
1096+
"""
1097+
if isinstance(config, dict):
1098+
config = types.AssessDatasetConfig(**config)
1099+
elif not config:
1100+
config = types.AssessDatasetConfig()
1101+
1102+
operation = self._assess_multimodal_dataset(
1103+
name=dataset_name,
1104+
batch_prediction_resource_usage_assessment_config=types.BatchPredictionResourceUsageAssessmentConfig(
1105+
model_name=model_name,
1106+
),
1107+
gemini_request_read_config=types.GeminiRequestReadConfig(
1108+
template_config=template_config,
1109+
),
1110+
config=config,
1111+
)
1112+
response = self._wait_for_operation(
1113+
operation=operation,
1114+
timeout_seconds=config.timeout,
1115+
)
1116+
result = response["batchPredictionResourceUsageAssessmentResult"]
1117+
return _datasets_utils.create_from_response(
1118+
types.BatchPredictionResourceUsageAssessmentResult, result
1119+
)
1120+
10571121

10581122
class AsyncDatasets(_api_module.BaseModule):
10591123

@@ -1875,3 +1939,67 @@ async def assess_tuning_resources(
18751939
types.TuningResourceUsageAssessmentResult,
18761940
response["tuningResourceUsageAssessmentResult"],
18771941
)
1942+
1943+
async def assess_batch_prediction_resources(
1944+
self,
1945+
*,
1946+
dataset_name: str,
1947+
model_name: str,
1948+
template_config: Optional[types.GeminiTemplateConfigOrDict] = None,
1949+
config: Optional[types.AssessDatasetConfigOrDict] = None,
1950+
) -> types.BatchPredictionResourceUsageAssessmentResult:
1951+
"""Assess the batch prediction resources required for a given model.
1952+
1953+
Args:
1954+
dataset_name:
1955+
Required. The name of the dataset to assess the batch prediction
1956+
resources.
1957+
model_name:
1958+
Required. The name of the model to assess the batch prediction
1959+
resources.
1960+
template_config:
1961+
Optional. The template config used to assemble the dataset
1962+
before assessing the batch prediction resources. If not provided,
1963+
the template config attached to the dataset will be used. Required
1964+
if no template config is attached to the dataset.
1965+
template_config:
1966+
Optional. The template config used to assemble the dataset
1967+
before assessing the batch prediction resources. If not provided, the
1968+
template config attached to the dataset will be used. Required
1969+
if no template config is attached to the dataset.
1970+
config:
1971+
Optional. A configuration for assessing the batch prediction
1972+
resources. If not provided, the default configuration will be
1973+
used.
1974+
1975+
Returns:
1976+
A types.BatchPredictionResourceUsageAssessmentResult object
1977+
representing the batch prediction resource usage assessment result.
1978+
It contains the following keys:
1979+
- token_count: The number of tokens in the dataset.
1980+
- audio_token_count: The number of audio tokens in the dataset.
1981+
1982+
"""
1983+
if isinstance(config, dict):
1984+
config = types.AssessDatasetConfig(**config)
1985+
elif not config:
1986+
config = types.AssessDatasetConfig()
1987+
1988+
operation = await self._assess_multimodal_dataset(
1989+
name=dataset_name,
1990+
batch_prediction_resource_usage_assessment_config=types.BatchPredictionResourceUsageAssessmentConfig(
1991+
model_name=model_name,
1992+
),
1993+
gemini_request_read_config=types.GeminiRequestReadConfig(
1994+
template_config=template_config,
1995+
),
1996+
config=config,
1997+
)
1998+
response = await self._wait_for_operation(
1999+
operation=operation,
2000+
timeout_seconds=config.timeout,
2001+
)
2002+
result = response["batchPredictionResourceUsageAssessmentResult"]
2003+
return _datasets_utils.create_from_response(
2004+
types.BatchPredictionResourceUsageAssessmentResult, result
2005+
)

0 commit comments

Comments
 (0)