Skip to content

Commit 1adf72b

Browse files
Ark-kuncopybara-github
authored andcommitted
fix: LLM - Fixed get_tuned_model for the future models that are not text-bison
Also fixed the `_endpoint_name` attribute of tuned models although it's not used. PiperOrigin-RevId: 551363402
1 parent 3bfd996 commit 1adf72b

File tree

1 file changed

+5
-8
lines changed

1 file changed

+5
-8
lines changed

vertexai/language_models/_language_models.py

+5-8
Original file line numberDiff line numberDiff line change
@@ -42,21 +42,17 @@ def _get_model_id_from_tuning_model_id(tuning_model_id: str) -> str:
4242
"""Gets the base model ID for the model ID labels used the tuned models.
4343
4444
Args:
45-
tuning_model_id: The model ID used in tuning
45+
tuning_model_id: The model ID used in tuning. E.g. `text-bison-001`
4646
4747
Returns:
4848
The publisher model ID
4949
5050
Raises:
5151
ValueError: If tuning model ID is unsupported
5252
"""
53-
if tuning_model_id.startswith("text-bison-"):
54-
return tuning_model_id.replace(
55-
"text-bison-", "publishers/google/models/text-bison@"
56-
)
57-
if "/" not in tuning_model_id:
58-
return "publishers/google/models/" + tuning_model_id
59-
return tuning_model_id
53+
model_name, _, version = tuning_model_id.rpartition("-")
54+
# "publishers/google/models/text-bison@001"
55+
return f"publishers/google/models/{model_name}@{version}"
6056

6157

6258
class _LanguageModel(_model_garden_models._ModelGardenModel):
@@ -203,6 +199,7 @@ def tune_model(
203199
tuned_model = job.result()
204200
# The UXR study attendees preferred to tune model in place
205201
self._endpoint = tuned_model._endpoint
202+
self._endpoint_name = tuned_model._endpoint_name
206203

207204

208205
@dataclasses.dataclass

0 commit comments

Comments
 (0)