google-cloud-pipeline-components 2.13.1__py3-none-any.whl → 2.14.1__py3-none-any.whl

google_cloud_pipeline_components/__init__.py

@@ -17,14 +17,13 @@ import warnings
 
 from google_cloud_pipeline_components.version import __version__
 
-if sys.version_info < (3, 8):
+if sys.version_info < (3, 9):
   warnings.warn(
       (
-          'Python 3.7 has reached end-of-life. Google Cloud Pipeline Components'
-          ' will drop support for Python 3.7 on April 23, 2024. To use new'
-          ' versions of the KFP SDK after that date, you will need to upgrade'
-          ' to Python >= 3.8. See https://devguide.python.org/versions/ for'
-          ' more details.'
+          ' Google Cloud Pipeline Components will drop support for Python 3.8'
+          ' on Oct 1, 2024. To use new versions of the GCPC SDK after that'
+          ' date, you will need to upgrade to Python >= 3.9. See'
+          ' https://devguide.python.org/versions/ for more details.'
       ),
       FutureWarning,
       stacklevel=2,

google_cloud_pipeline_components/_implementation/llm/deployment_graph.py

@@ -34,10 +34,13 @@ PipelineOutput = NamedTuple(
 def pipeline(
     output_adapter_path: str,
     large_model_reference: str,
+    policy_model_reference: str,
     model_display_name: Optional[str] = None,
     deploy_model: bool = True,
+    upload_model: bool = True,
     encryption_spec_key_name: str = '',
     upload_location: str = _placeholders.LOCATION_PLACEHOLDER,
+    regional_endpoint: str = '',
 ) -> PipelineOutput:
   # fmt: off
   """Uploads a tuned language model and (optionally) deploys it to an endpoint.
@@ -45,62 +48,37 @@ def pipeline(
   Args:
     output_adapter_path: Path to the trained model adapter if LoRA tuning was used.
     large_model_reference: Name of the base model. Supported values are `text-bison@001`, `t5-small`, `t5-large`, `t5-xl` and `t5-xxl`. `text-bison@001` and `t5-small` are supported in `us-central1` and `europe-west4`. `t5-large`, `t5-xl` and `t5-xxl` are only supported in `europe-west4`.
+    policy_model_reference: The name of the model for deployment. The name should be in capitalized snake case format.
     model_display_name: Name of the fine-tuned model shown in the Model Registry. If not provided, a default name will be created.
     deploy_model: Whether to deploy the model to an endpoint in `us-central1`. Default is True.
     encryption_spec_key_name: Customer-managed encryption key. If this is set, then all resources created by the CustomJob will be encrypted with the provided encryption key. Note that this is not supported for TPU at the moment.
     upload_location: Region to upload and deploy the model to. Default is the location used to run the pipeline components.
+    regional_endpoint: Regional endpoint to upload the model.
 
   Returns:
     model_resource_name: Path to the model uploaded to the Model Registry. This will be an empty string if the model was not deployed.
     endpoint_resource_name: Path the Online Prediction Endpoint. This will be an empty string if the model was not deployed.
   """
   # fmt: on
-  regional_endpoint = function_based.resolve_regional_endpoint(
-      upload_location=upload_location
-  ).set_display_name('Resolve Regional Endpoint')
-
-  display_name = (
-      function_based.resolve_model_display_name(
-          large_model_reference=large_model_reference,
-          model_display_name=model_display_name,
-      )
-      .set_caching_options(False)
-      .set_display_name('Resolve Model Display Name')
-  )
-
-  reference_model_metadata = function_based.resolve_reference_model_metadata(
-      large_model_reference=large_model_reference,
-  ).set_display_name('Resolve Model Metadata')
-
-  upload_model = function_based.resolve_upload_model(
-      large_model_reference=reference_model_metadata.outputs[
-          'large_model_reference'
-      ]
-  ).set_display_name('Resolve Upload Model')
   upload_task = upload_llm_model.refined_upload_llm_model(
       project=_placeholders.PROJECT_ID_PLACEHOLDER,
       location=upload_location,
-      regional_endpoint=regional_endpoint.output,
+      regional_endpoint=regional_endpoint,
       artifact_uri=output_adapter_path,
-      model_display_name=display_name.output,
+      model_display_name=model_display_name,
       model_reference_name=large_model_reference,
-      upload_model=upload_model.output,
+      upload_model=upload_model,
       encryption_spec_key_name=encryption_spec_key_name,
       tune_type='rlhf',
   ).set_display_name('Upload Model')
-  deploy_model = function_based.resolve_deploy_model(
-      deploy_model=deploy_model,
-      large_model_reference=reference_model_metadata.outputs[
-          'large_model_reference'
-      ],
-  ).set_display_name('Resolve Deploy Model')
+
   deploy_task = deploy_llm_model.deploy_llm_model(
       project=_placeholders.PROJECT_ID_PLACEHOLDER,
       location=upload_location,
       model_resource_name=upload_task.outputs['model_resource_name'],
-      display_name=display_name.output,
-      regional_endpoint=regional_endpoint.output,
-      deploy_model=deploy_model.output,
+      display_name=model_display_name,
+      regional_endpoint=regional_endpoint,
+      deploy_model=deploy_model,
       encryption_spec_key_name=encryption_spec_key_name,
   ).set_display_name('Deploy Model')
   return PipelineOutput(

google_cloud_pipeline_components/_implementation/llm/env.py

@@ -19,7 +19,7 @@ from google_cloud_pipeline_components._implementation.llm.generated import refin
 
 
 def get_private_image_tag() -> str:
-  return os.getenv('PRIVATE_IMAGE_TAG') or '20240330_0352_RC00'
+  return os.getenv('PRIVATE_IMAGE_TAG') or refined_image_versions.IMAGE_TAG
 
 
 def get_autosxs_image_tag() -> str:

google_cloud_pipeline_components/_implementation/llm/function_based.py

@@ -231,8 +231,8 @@ def resolve_reference_model_metadata(
               'gs://vertex-llm-restricted/cloud-llm-restricted/checkpoints/'
               'safe_flan_t5/xxl/v1/checkpoint_1190000/'
           ),
-          reward_model_reference='T5_XL',
-          reward_model_path='gs://t5-data/pretrained_models/t5x/t5_1_1_xl',
+          reward_model_reference='T5_XXL',
+          reward_model_path='gs://t5-data/pretrained_models/t5x/t5_1_1_xxl',
           is_supported=True,
       ),
       'palm-tiny': reference_model_metadata(
@@ -265,8 +265,10 @@ def resolve_reference_model_metadata(
           reference_model_path=(
               'gs://vertex-rlhf-restricted/pretrained_models/palm/t5x_bison/'
           ),
-          reward_model_reference='OTTER',
-          reward_model_path='gs://vertex-rlhf-restricted/pretrained_models/palm/t5x_otter_pretrain/',
+          reward_model_reference='BISON',
+          reward_model_path=(
+              'gs://vertex-rlhf-restricted/pretrained_models/palm/t5x_bison/'
+          ),
           is_supported=False,  # Deprecated: Use text-bision@001 instead.
       ),
       'text-bison@001': reference_model_metadata(
@@ -274,8 +276,10 @@ def resolve_reference_model_metadata(
           reference_model_path=(
               'gs://vertex-rlhf-restricted/pretrained_models/palm/t5x_bison/'
           ),
-          reward_model_reference='OTTER',
-          reward_model_path='gs://vertex-rlhf-restricted/pretrained_models/palm/t5x_otter_pretrain/',
+          reward_model_reference='BISON',
+          reward_model_path=(
+              'gs://vertex-rlhf-restricted/pretrained_models/palm/t5x_bison/'
+          ),
           is_supported=True,
       ),
       'text-bison@002': reference_model_metadata(
@@ -292,8 +296,10 @@ def resolve_reference_model_metadata(
           reference_model_path=(
               'gs://vertex-rlhf-restricted/pretrained_models/palm/t5x_bison/'
           ),
-          reward_model_reference='OTTER',
-          reward_model_path='gs://vertex-rlhf-restricted/pretrained_models/palm/t5x_otter_pretrain/',
+          reward_model_reference='BISON',
+          reward_model_path=(
+              'gs://vertex-rlhf-restricted/pretrained_models/palm/t5x_bison/'
+          ),
           is_supported=True,
       ),
       'elephant': reference_model_metadata(
@@ -372,46 +378,6 @@ def convert_to_delimited_string(items: List[str], delimiter: str = ',') -> str:
   return delimiter.join(items)
 
 
-@dsl.component(base_image=_image.GCPC_IMAGE_TAG, install_kfp_package=False)
-def generate_default_instruction(
-    task: str,
-    target_sequence_length: int,
-    instruction_override: str = '',
-) -> str:
-  """Generates a default instruction if no override is provided."""
-  if instruction_override:
-    return instruction_override
-  task = task.lower()
-  if task == 'summarization':
-    return f'Summarize in less than {target_sequence_length} words.'
-
-  elif task == 'question_answer':
-    return f'Answer the question in less than {target_sequence_length} words.'
-
-  else:
-    raise ValueError(
-        f'Task not recognized: {task}. Supported tasks are: "summarization",'
-        ' "question_answer".'
-    )
-
-
-@dsl.component(base_image=_image.GCPC_IMAGE_TAG, install_kfp_package=False)
-def resolve_upload_location(upload_location: Optional[str] = None) -> str:
-  """Gets the region to upload the model.
-
-  Args:
-    upload_location: User-specified region to upload the model to.
-
-  Returns:
-    Where to upload the model. If no location is specified, the model will be
-    uploaded to the region where the pipeline is running.
-  """
-  # pylint: disable=g-import-not-at-top
-  import os
-  # pylint: enable=g-import-not-at-top
-  return upload_location or os.environ['CLOUD_ML_REGION']
-
-
 @dsl.component(base_image=_image.GCPC_IMAGE_TAG, install_kfp_package=False)
 def resolve_regional_endpoint(upload_location: str) -> str:
   """Gets the regional endpoint used to upload a model to the registry.

google_cloud_pipeline_components/_implementation/llm/generated/refined_image_versions.py

@@ -17,4 +17,4 @@
 DO NOT EDIT - This file is generated, manual changes will be overridden.
 """
 
-IMAGE_TAG = '20240327_1338'
+IMAGE_TAG = '20240506_1707'

google_cloud_pipeline_components/_implementation/llm/infer_preprocessor.py

@@ -0,0 +1,109 @@
+# Copyright 2024 The Kubeflow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Component that preprocesses inputs for infer pipeline."""
+
+from google_cloud_pipeline_components import _placeholders
+from google_cloud_pipeline_components import utils as gcpc_utils
+from google_cloud_pipeline_components._implementation.llm import utils
+from kfp import dsl
+
+
+@dsl.container_component
+def infer_preprocessor(
+    large_model_reference: str,
+    accelerator_type: str,
+    use_test_spec: bool,
+    project: str,
+    location: str,
+    artifact_registry: str,
+    tag: str,
+    gcp_resources: dsl.OutputPath(str),  # pytype: disable=invalid-annotation
+    metadata_large_model_reference: dsl.OutputPath(str),  # pytype: disable=invalid-annotation
+    metadata_reference_model_path: dsl.OutputPath(str),  # pytype: disable=invalid-annotation
+    metadata_reward_model_reference: dsl.OutputPath(str),  # pytype: disable=invalid-annotation
+    metadata_reward_model_path: dsl.OutputPath(str),  # pytype: disable=invalid-annotation
+    metadata_machine_type: dsl.OutputPath(str),  # pytype: disable=invalid-annotation
+    metadata_tuning_location: dsl.OutputPath(str),  # pytype: disable=invalid-annotation
+    metadata_accelerator_type: dsl.OutputPath(str),  # pytype: disable=invalid-annotation
+    metadata_accelerator_count: dsl.OutputPath(int),  # pytype: disable=invalid-annotation
+    metadata_instruction: dsl.OutputPath(str),  # pytype: disable=invalid-annotation
+    metadata_refined_image_uri: dsl.OutputPath(str),  # pytype: disable=invalid-annotation
+    use_experimental_image: bool = False,
+    input_reference_model_path: str = '',
+    instruction: str = '',
+    image_uri: str = utils.get_default_image_uri('refined_cpu', ''),
+) -> dsl.ContainerSpec:  # pylint: disable=g-doc-args
+  # fmt: off
+  """Preprocess infer pipeline inputs.
+
+  Args:
+    large_model_reference: The model for fine tuning.
+    accelerator_type: Specific accelerator type for the job.
+    use_test_spec: Whether to use a lower resource machine for testing.
+    project: Project that contains the artifact registry.
+    location: Region that contains the artifact registry.
+    artifact_registry: Registry that contains Docker images.
+    tag: Image tag.
+    use_experimental_image:  Whether to use refined experimental image.
+    input_reference_model_path: The model checkpoint path for the reference model
+    instruction: The instruction to let the model know what task it needs to perform.
+    image_uri: Docker image URI to use for the custom job.
+
+  Returns:
+    gcp_resources: GCP resources that can be used to track the custom job.
+    metadata_large_model_reference: The base model for fine tuning. The name should be in capitalized snake case format.
+    metadata_reference_model_path: The model checkpoint path for the reinforcer model
+    metadata_reward_model_reference:  The base model for training reward model. The name should be in capitalized snake case format.
+    metadata_reward_model_path: The model checkpoint path for the reward model.
+    metadata_machine_type: The type of the machine to provision for the custom job.
+    metadata_tuning_location: The GCP region to run the custom job.
+    metadata_accelerator_type: Specific accelerator type for the custom job.
+    metadata_accelerator_count: The number of accelerator.
+    metadata_instruction: The instruction to let the model know what task it needs to perform.
+    metadata_refined_image_uri: Docker image URI to use for the custom job.
+  """
+  # fmt: on
+  return gcpc_utils.build_serverless_customjob_container_spec(
+      project=_placeholders.PROJECT_ID_PLACEHOLDER,
+      location=_placeholders.LOCATION_PLACEHOLDER,
+      custom_job_payload=utils.build_payload(
+          display_name='infer_preprocessor',
+          machine_type='n1-standard-4',
+          image_uri=image_uri,
+          args=[
+              '--app_name=infer_preprocessor',
+              f'--large_model_reference={large_model_reference}',
+              f'--input_reference_model_path={input_reference_model_path}',
+              f'--accelerator_type={accelerator_type}',
+              f'--use_test_spec={use_test_spec}',
+              f'--project={project}',
+              f'--location={location}',
+              f'--artifact_registry={artifact_registry}',
+              f'--tag={tag}',
+              f'--use_experimental_image={use_experimental_image}',
+              f'--instruction={instruction}',
+              f'--metadata_large_model_reference_path={metadata_large_model_reference}',
+              f'--metadata_reference_model_path_path={metadata_reference_model_path}',
+              f'--metadata_reward_model_reference_path={metadata_reward_model_reference}',
+              f'--metadata_reward_model_path_path={metadata_reward_model_path}',
+              f'--metadata_machine_type_path={metadata_machine_type}',
+              f'--metadata_tuning_location_path={metadata_tuning_location}',
+              f'--metadata_accelerator_type_path={metadata_accelerator_type}',
+              f'--metadata_accelerator_count_path={metadata_accelerator_count}',
+              f'--metadata_instruction_path={metadata_instruction}',
+              f'--metadata_refined_image_uri_path={metadata_refined_image_uri}',
+          ],
+      ),
+      gcp_resources=gcp_resources,
+  )

google_cloud_pipeline_components/_implementation/llm/online_evaluation_pairwise.py

@@ -52,6 +52,7 @@ def online_evaluation_pairwise(
     project: str = _placeholders.PROJECT_ID_PLACEHOLDER,
     location: str = _placeholders.LOCATION_PLACEHOLDER,
     encryption_spec_key_name: str = '',
+    autorater_prompt_parameters: Dict[str, Dict[str, str]] = {},
 ) -> dsl.ContainerSpec:  # pylint: disable=g-doc-args
   """Evaluate two models using an autorater.
 
@@ -73,6 +74,8 @@ def online_evaluation_pairwise(
     encryption_spec_key_name: Customer-managed encryption key options. If this
       is set, then all resources created by the component will be encrypted with
       the provided encryption key.
+    autorater_prompt_parameters: Map of autorater prompt template parameters to
+      columns or templates.
 
   Returns:
     judgments: Individual judgments used to calculate the win rates.
@@ -112,6 +115,11 @@ def online_evaluation_pairwise(
               '--executor_input={{$.json_escape[1]}}',
               f'--kms_key_name={encryption_spec_key_name}',
               f'--metadata_path={metadata}',
+              (
+                  '--autorater_prompt_parameters='
+                  "{{$.inputs.parameters['autorater_prompt_parameters']"
+                  '.json_escape[0]}}'
+              ),
           ],
           encryption_spec_key_name=encryption_spec_key_name,
       ),

google_cloud_pipeline_components/_implementation/llm/reinforcement_learning_graph.py

@@ -41,6 +41,14 @@ def pipeline(
     input_reward_adapter_path: str,
     input_preference_dataset_path: str,
     large_model_reference: str,
+    reward_model_reference: str,
+    policy_model_reference: str,
+    policy_model_path: str,
+    machine_type: str,
+    tuning_location: str,
+    accelerator_type: str,
+    accelerator_count: int,
+    rl_image_uri: str,
     prompt_sequence_length: int = 512,
     target_sequence_length: int = 64,
     lora_dim: int = 1,
@@ -51,10 +59,10 @@ def pipeline(
     kl_coeff: float = 0.1,
     instruction: Optional[str] = None,
     project: str = _placeholders.PROJECT_ID_PLACEHOLDER,
-    accelerator_type: str = 'GPU',
     location: str = _placeholders.LOCATION_PLACEHOLDER,
     tensorboard_resource_id: str = '',
     encryption_spec_key_name: str = '',
+    num_microbatches: int = 0,
 ) -> PipelineOutput:
   # fmt: off
   """Trains a reward model.
@@ -64,6 +72,14 @@ def pipeline(
     input_reward_adapter_path: Path to the reward LoRA adapter to use during reinforcement learning.
     input_preference_dataset_path: Path to preference dataset used by the reward model.
     large_model_reference: Name of the base model. Supported values are `text-bison@001`, `t5-small`, `t5-large`, `t5-xl` and `t5-xxl`. `text-bison@001` and `t5-small` are supported in `us-central1` and `europe-west4`. `t5-large`, `t5-xl` and `t5-xxl` are only supported in `europe-west4`.
+    reward_model_reference: Name of the reward model. The name should be in capitalized snake case format.
+    policy_model_reference: Name of the policy model. The name should be in capitalized snake case format.
+    policy_model_path: The model checkpoint path to the reinforcer model.
+    machine_type: The type of the machine to provision for the custom job. Must be a valid GCE instance type and compatible with the accelerator type.
+    tuning_location: The GCP region to run the custom job.
+    accelerator_type: Specific accelerator type for the custom job.
+    accelerator_count: The number of accelerator.
+    rl_image_uri: Docker image URI to use for the reinforcement learning training job.
     prompt_sequence_length: Maximum tokenized sequence length for input text. Higher values increase memory overhead. This value should be at most 8192. Default value is 512.
     target_sequence_length: Maximum tokenized sequence length for target text. Higher values increase memory overhead. This value should be at most 1024. Default value is 64.
     lora_dim: The rank of the LoRA adapter. If >0, then use LoRA-tuning. If =0, then use full-tuning. Default is 1.
@@ -74,7 +90,6 @@ def pipeline(
     kl_coeff: Coefficient for KL penalty. This regularizes the policy model and penalizes if it diverges from its initial distribution. If set to 0, the reference language model is not loaded into memory. Default value is 0.1.
     instruction: This field lets the model know what task it needs to perform. Base models have been trained over a large set of varied instructions. You can give a simple and intuitive description of the task and the model will follow it, e.g. "Classify this movie review as positive or negative" or "Translate this sentence to Danish". Do not specify this if your dataset already prepends the instruction to the inputs field.
     project: Project used to run custom jobs. If not specified the project used to run the pipeline will be used.
-    accelerator_type: One of 'TPU' or 'GPU'. If 'TPU' is specified, tuning components run in europe-west4. Otherwise tuning components run in us-central1 on GPUs. Default is 'GPU'.
     location: Location used to run non-tuning components, i.e. components that do not require accelerators. If not specified the location used to run the pipeline will be used.
     tensorboard_resource_id: Optional tensorboard resource id in format `projects/{project_number}/locations/{location}/tensorboards/{tensorboard_id}`. If provided, tensorboard metrics will be uploaded to this location.
     encryption_spec_key_name: Customer-managed encryption key. If this is set, then all resources created by the CustomJob will be encrypted with the provided encryption key. Note that this is not supported for TPU at the moment.
@@ -85,14 +100,6 @@ def pipeline(
   """
   # fmt: on
   prompt_column = 'input_text'
-  machine_spec = function_based.resolve_machine_spec(
-      accelerator_type=accelerator_type,
-      use_test_spec=env.get_use_test_machine_spec(),
-  ).set_display_name('Resolve Machine Spec')
-
-  reference_model_metadata = function_based.resolve_reference_model_metadata(
-      large_model_reference=large_model_reference,
-  ).set_display_name('Resolve Model Metadata')
 
   processed_dataset = preprocess_chat_dataset.preprocess_chat_dataset(
       large_model_reference=large_model_reference,
@@ -109,30 +116,18 @@ def pipeline(
           # Target field name does not matter because this field is not used.
           targets_field_name='non_existent_targets_field_name',
           output_split_name=env.TRAIN_SPLIT,
-          large_model_reference=reference_model_metadata.outputs[
-              'large_model_reference'
-          ],
+          large_model_reference=policy_model_reference,
           instruction=instruction,
           encryption_spec_key_name=encryption_spec_key_name,
       )
       .set_display_name('Import Prompt Dataset')
       .set_caching_options(False)
   )
-  rl_image_uri = function_based.resolve_private_refined_image_uri(
-      accelerator_type=machine_spec.outputs['accelerator_type'],
-  ).set_display_name('Resolve Reinforcer Image URI')
-  num_microbatches = function_based.resolve_num_microbatches(
-      large_model_reference=reference_model_metadata.outputs[
-          'large_model_reference'
-      ]
-  ).set_display_name('Resolve Number of Microbatches')
   rl_model = (
       reinforcer.reinforcer(
           project=project,
-          location=machine_spec.outputs['tuning_location'],
-          input_reference_model_path=reference_model_metadata.outputs[
-              'reference_model_path'
-          ],
+          location=tuning_location,
+          input_reference_model_path=policy_model_path,
           input_reward_model_path=input_reward_model_path,
           input_reward_adapter_path=input_reward_adapter_path,
           input_dataset_path=prompt_dataset_importer.outputs[
@@ -140,16 +135,12 @@ def pipeline(
           ],
           input_preference_dataset_path=input_preference_dataset_path,
           train_steps=reinforcement_learning_train_steps,
-          accelerator_type=machine_spec.outputs['accelerator_type'],
-          accelerator_count=machine_spec.outputs['accelerator_count'],
-          large_model_reference=reference_model_metadata.outputs[
-              'large_model_reference'
-          ],
-          reward_model_reference=reference_model_metadata.outputs[
-              'reward_model_reference'
-          ],
-          machine_type=machine_spec.outputs['machine_type'],
-          image_uri=rl_image_uri.output,
+          accelerator_type=accelerator_type,
+          accelerator_count=accelerator_count,
+          large_model_reference=policy_model_reference,
+          reward_model_reference=reward_model_reference,
+          machine_type=machine_type,
+          image_uri=rl_image_uri,
           inputs_sequence_length=prompt_sequence_length,
           targets_sequence_length=target_sequence_length,
           batch_size=batch_size,
@@ -157,7 +148,7 @@ def pipeline(
           kl_coeff=kl_coeff,
           lora_dim=lora_dim,
           reward_lora_dim=reward_lora_dim,
-          num_microbatches=num_microbatches.output,
+          num_microbatches=num_microbatches,
           encryption_spec_key_name=encryption_spec_key_name,
           tensorboard_resource_id=tensorboard_resource_id,
       )