Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

feat: Expose new text embedding tuning parameters in GA namespace. #3830

Merged
merged 1 commit into from
Jun 24, 2024
Merged

feat: Expose new text embedding tuning parameters in GA namespace. #3830

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 2 additions & 8 deletions tests/unit/aiplatform/test_language_models.py
View file
Open in desktop
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
Open in desktop
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