def extract_embeddings()

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

• 19 lines of code
• 2 McCabe index (conditional complexity)

def extract_embeddings(
  project_id: Parameter[str],
  bq_dataset: Parameter[str],
  bq_model: InputArtifact[BQModel],
  item_embeddings: OutputArtifact[Dataset]):
  
  embedding_model_name = bq_model.get_string_custom_property('bq_model_name')
  stored_proc = f'{bq_dataset}.sp_ExractEmbeddings'
  query = f'''
      CALL {stored_proc}();
  '''
  result_table = 'item_embeddings'

  logging.info(f'Extracting item embedding from: {bq_dataset}.{embedding_model_name}')
  logging.info(f'Starting exporting embeddings...')
  
  client = bigquery.Client(project=project_id)
  query_job = client.query(query)
  query_job.result() # Wait for the job to complete
  
  logging.info(f'Embeddings export completed. Output in {bq_dataset}.{result_table}')
  
  # Write the location of the output table to metadata.
  item_embeddings.set_string_custom_property('bq_dataset', bq_dataset)
  item_embeddings.set_string_custom_property('bq_result_table', result_table)