def main()

in tensorflow_hub/tools/make_nearest_neighbour_index/make_nearest_neighbour_index.py [0:0]


def main(args):
  """Entry point main function."""

  operation = validate_args(args)
  print("Selected operation: {}".format(operation))

  if operation == "generate":
    print("Generating embeddings...")
    generator.run(FLAGS)
    print("Embedding generation completed.")

  elif operation == "build":
    print("Building ANN index...")
    builder.run(FLAGS)
    print("Building ANN index completed.")

  elif operation == "e2e":
    print("Generating embeddings and building ANN index...")
    generator.run(FLAGS)
    print("Embedding generation completed.")
    if FLAGS.projected_dim:
      FLAGS.dimensions = FLAGS.projected_dim

    builder.run(FLAGS)
    print("Building ANN index completed.")

  else:
    print("Querying the ANN index...")
    similarity_finder = finder.load(FLAGS)
    num_matches = FLAGS.num_matches
    while True:
      print("Enter your query: ", end="")
      query = str(input())
      similar_items = similarity_finder.find_similar_items(query, num_matches)
      print("Results:")
      print("=========")
      for item in similar_items:
        print(item)