def train_item_matching_model()

in retail/recommendation-system/bqml-scann/tfx_pipeline/bq_components.py [0:0]

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

def train_item_matching_model(
  project_id: Parameter[str],
  bq_dataset: Parameter[str],
  dimensions: Parameter[int],
  item_cooc: InputArtifact[Dataset],
  bq_model: OutputArtifact[BQModel]):
    
  item_cooc_table = item_cooc.get_string_custom_property('bq_result_table')
  stored_proc = f'{bq_dataset}.sp_TrainItemMatchingModel'
  query = f'''
    DECLARE dimensions INT64 DEFAULT {dimensions};
    CALL {stored_proc}(dimensions);
  '''
  model_name = 'item_matching_model'
  
  logging.info(f'Using item co-occurrence table: {bq_dataset}.{item_cooc_table}')
  logging.info(f'Starting training of the model...')
    
  client = bigquery.Client(project=project_id)
  query_job = client.query(query)
  query_job.result()
  
  logging.info(f'Model training completed. Output in {bq_dataset}.{model_name}.')
  
  # Write the location of the model to metadata.  
  bq_model.set_string_custom_property('bq_dataset', bq_dataset)
  bq_model.set_string_custom_property('bq_model_name', model_name)