def build_index()

in retail/recommendation-system/bqml-scann/index_builder/builder/indexer.py [0:0]


def build_index(embeddings, num_leaves):
  
  data_size = embeddings.shape[0] 
  if not num_leaves:
    num_leaves = int(math.sqrt(data_size))
    
  print('Start building the ScaNN index...')
  scann_builder = scann.scann_ops.builder(embeddings, NUM_NEIGHBOURS, METRIC)
  scann_builder = scann_builder.tree(
    num_leaves=num_leaves, 
    num_leaves_to_search=NUM_LEAVES_TO_SEARCH, 
    training_sample_size=data_size)
  scann_builder = scann_builder.score_ah(
    DIMENSIONS_PER_BLOCK, 
    anisotropic_quantization_threshold=ANISOTROPIC_QUANTIZATION_THRESHOLD)
  scann_builder = scann_builder.reorder(REORDER_NUM_NEIGHBOURS)
  scann_index = scann_builder.build()
  print('ScaNN index is built.')
  
  return scann_index