Skip to content

Commit

Permalink
feat: Expose new text embedding tuning parameters in GA namespace.
Browse files Browse the repository at this point in the history 
PiperOrigin-RevId: 646160221
  • Loading branch information
@vertex-sdk-bot @copybara-github
vertex-sdk-bot authored and copybara-github committed Jun 24, 2024
1 parent d4b0091 commit 249a5fa
Show file tree
Hide file tree
Showing 2 changed files with 4 additions and 86 deletions.
10 changes: 2 additions & 8 deletions tests/unit/aiplatform/test_language_models.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2290,11 +2290,10 @@ def test_text_generation_response_repr(self):
indirect=True,
)
@pytest.mark.parametrize(
"base_model_version_id,use_preview_module,tune_args,expected_pipeline_args",
"base_model_version_id,tune_args,expected_pipeline_args",
[ # Do not pass any optional parameters.
(
"textembedding-gecko@003",
False,
dict(
training_data="gs://bucket/training.tsv",
corpus_data="gs://bucket/corpus.jsonl",
Expand All @@ -2311,7 +2310,6 @@ def test_text_generation_response_repr(self):
# Pass all optional parameters.
(
"text-multilingual-embedding-002",
True,
dict(
training_data="gs://bucket/training.tsv",
corpus_data="gs://bucket/corpus.jsonl",
Expand Down Expand Up @@ -2364,7 +2362,6 @@ def test_tune_text_embedding_model(
tune_args,
expected_pipeline_args,
base_model_version_id,
use_preview_module,
):
"""Tests tuning the text embedding model."""
aiplatform.init(
Expand All @@ -2379,10 +2376,7 @@ def test_tune_text_embedding_model(
_TEXT_GECKO_PUBLISHER_MODEL_DICT
),
):
language_models_module = (
preview_language_models if use_preview_module else language_models
)
model = language_models_module.TextEmbeddingModel.from_pretrained(
model = language_models.TextEmbeddingModel.from_pretrained(
base_model_version_id
)
tuning_result = model.tune_model(**tune_args)
Expand Down
80 changes: 2 additions & 78 deletions vertexai/language_models/_language_models.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2268,7 +2268,7 @@ def tune_model(
```
tuning_job = model.tune_model(...)
... do some other work
tuned_model = tuning_job.get_tuned_model() # Blocks until tuning is complete
tuned_model = tuning_job.deploy_tuned_model() # Blocks until tuning is complete
Args:
training_data: URI pointing to training data in TSV format.
Expand Down Expand Up @@ -2414,83 +2414,7 @@ def deploy_tuned_model(


class _TunableTextEmbeddingModelMixin(_PreviewTunableTextEmbeddingModelMixin):
def tune_model(
self,
*,
training_data: Optional[str] = None,
corpus_data: Optional[str] = None,
queries_data: Optional[str] = None,
test_data: Optional[str] = None,
validation_data: Optional[str] = None,
batch_size: Optional[int] = None,
train_steps: Optional[int] = None,
tuned_model_location: Optional[str] = None,
model_display_name: Optional[str] = None,
task_type: Optional[str] = None,
machine_type: Optional[str] = None,
accelerator: Optional[str] = None,
accelerator_count: Optional[int] = None,
) -> "_TextEmbeddingModelTuningJob":
"""Tunes a model based on training data.
This method launches and returns an asynchronous model tuning job.
Usage:
```
tuning_job = model.tune_model(...)
... do some other work
tuned_model = tuning_job.get_tuned_model() # Blocks until tuning is complete
Args:
training_data: URI pointing to training data in TSV format.
corpus_data: URI pointing to data in JSON lines format.
queries_data: URI pointing to data in JSON lines format.
test_data: URI pointing to data in TSV format.
validation_data: URI pointing to data in TSV format.
batch_size: The training batch size.
train_steps: The number of steps to perform model tuning. Must
be greater than 30.
tuned_model_location: GCP location where the tuned model should be deployed.
model_display_name: Custom display name for the tuned model.
task_type: The task type expected to be used during inference.
Valid values are `DEFAULT`, `RETRIEVAL_QUERY`, `RETRIEVAL_DOCUMENT`,
`SEMANTIC_SIMILARITY`, `CLASSIFICATION`, `CLUSTERING`,
`FACT_VERIFICATION`, and `QUESTION_ANSWERING`.
machine_type: The machine type to use for training. For information
about selecting the machine type that matches the accelerator
type and count you have selected, see
https://cloud.google.com/compute/docs/gpus.
accelerator: The accelerator type to use for tuning, for example
`NVIDIA_TESLA_V100`. For possible values, see
https://cloud.google.com/vertex-ai/generative-ai/docs/models/tune-embeddings#using-accelerators.
accelerator_count: The number of accelerators to use when training.
Using a greater number of accelerators may make training faster,
but has no effect on quality.
Returns:
A `LanguageModelTuningJob` object that represents the tuning job.
Calling `job.result()` blocks until the tuning is complete and
returns a `LanguageModel` object.
Raises:
ValueError: If the provided parameter combinations or values are not
supported.
RuntimeError: If the model does not support tuning
"""

return super().tune_model(
training_data=training_data,
corpus_data=corpus_data,
queries_data=queries_data,
test_data=test_data,
validation_data=validation_data,
task_type=task_type,
batch_size=batch_size,
train_steps=train_steps,
tuned_model_location=tuned_model_location,
model_display_name=model_display_name,
machine_type=machine_type,
accelerator=accelerator,
accelerator_count=accelerator_count,
)
pass


class TextEmbeddingModel(_TextEmbeddingModel, _TunableTextEmbeddingModelMixin):
Expand Down

0 comments on commit 249a5fa

Please sign in to comment.