def _predict()

in backend/time-series-forecasting/training_methods/automl_training_method.py [0:0]

• 18 lines of code
• 1 McCabe index (conditional complexity)

    def _predict(self, model_name: str, bigquery_source: str) -> str:
        """This function runs the batch prediction job

        Args:
            model_name (str): The model used for batch prediciton.
            bigquery_source (str): The BigQuery source URI for batch prediction.

        Returns:
            str: The table id of batch prediction results.
        """

        client = bigquery.Client()
        project_id = client.project

        model = aiplatform.Model(model_name=model_name)

        job = model.batch_predict(
            job_display_name=f"automl_forecasting_{utils.generate_uuid()}",
            bigquery_source=bigquery_source,
            instances_format="bigquery",
            bigquery_destination_prefix=f"bq://{project_id}",
            predictions_format="bigquery",
            generate_explanation=True,
            sync=True,
        )

        output_dataset = job.output_info.bigquery_output_dataset
        output_dataset = output_dataset.replace("bq://", "")
        output_table = job.output_info.bigquery_output_table
        bq_output_table_id = f"{output_dataset}.{output_table}"

        return bq_output_table_id