Adds a new --lsh_bits flag to nsl.tools.build_graph().

If --lsh_bits > 0, all the input points are first bucketed into
at most 2^n buckets, where 'n' is the --lsh_bits value. Each of
the buckets is then considered independently, vastly reducing
the number of pairs of points that need to be compared for
similarity. Thus, the running time of the program can be reduced
dramatically when the number of input points is large, at the
potential expense of missing some graph edges.

This also introduces a new configuration class and a new top-level function:
  nsl.configs.GraphBuilderConfig: graph building configuration object.
  nsl.tools.build_graph_from_config(...): parameterized by GraphBuilderConfig.

PiperOrigin-RevId: 318376805

Author: Neural-Link Team

neural_structured_learning/configs/__init__.py

@@ -10,6 +10,7 @@
 from neural_structured_learning.configs.configs import DEFAULT_DISTANCE_PARAMS
 from neural_structured_learning.configs.configs import DistanceConfig
 from neural_structured_learning.configs.configs import DistanceType
+from neural_structured_learning.configs.configs import GraphBuilderConfig
 from neural_structured_learning.configs.configs import GraphNeighborConfig
 from neural_structured_learning.configs.configs import GraphRegConfig
 from neural_structured_learning.configs.configs import IntegrationConfig
@@ -31,6 +32,7 @@
     'DEFAULT_DISTANCE_PARAMS',
     'DistanceConfig',
     'DistanceType',
+    'GraphBuilderConfig',
     'GraphNeighborConfig',
     'GraphRegConfig',
     'IntegrationConfig',

neural_structured_learning/configs/configs.py

@@ -17,8 +17,9 @@
 from __future__ import division
 from __future__ import print_function
 
-import attr
 import enum
+
+import attr
 import tensorflow as tf
 
 
@@ -291,6 +292,36 @@ class VirtualAdvConfig(object):
   approx_difference = attr.ib(default=1e-6)
 
 
+@attr.s
+class GraphBuilderConfig(object):
+  """Encapsulates configuration parameters for building a graph.
+
+  For more information, see `nsl.tools.build_graph_from_config`.
+
+  Attributes:
+    id_feature_name: The name of the feature in the input `tf.train.Example`
+      objects representing the ID of examples.
+    embedding_feature_name: The name of the feature in the input
+      `tf.train.Example` objects representing the embedding of examples.
+    similarity_threshold: Threshold used to determine which edges to retain in
+      the resulting graph.
+    lsh_bits: Determines the maximum number of LSH buckets into which input data
+      points will be bucketed by the graph builder. See the
+      `nsl.tools.build_graph_from_config` documentation for details. This
+      defaults to 0, in which case all pairs of inputs will be compared,
+      probably resulting in slow running times on larger input sets.
+    random_seed: Value used to seed the random number generator used to perform
+      randomized LSH bucketing of the inputs when `lsh_bits > 0`. By default,
+      the generator will be initialized randomly, but setting this to any
+      integer will initialize it deterministically.
+  """
+  id_feature_name = attr.ib(default='id')
+  embedding_feature_name = attr.ib(default='embedding')
+  similarity_threshold = attr.ib(default=0.8)
+  lsh_bits = attr.ib(default=0)
+  random_seed = attr.ib(default=None)
+
+
 @attr.s
 class GraphNeighborConfig(object):
   """Specifies neighbor attributes for graph regularization.

neural_structured_learning/tools/BUILD

@@ -63,6 +63,7 @@ py_library(
         # package absl:app
         # package absl/flags
         # package absl/logging
+        "//neural_structured_learning/configs",
         # package numpy
         # package six
         # package tensorflow

neural_structured_learning/tools/__init__.py

@@ -1,18 +1,28 @@
 """Tools and APIs for preparing data for Neural Structured Learning.
 
 In addition to the functions exported here, two of the modules can be invoked
-from the command-line as follows:
+from the command-line.
+
+Sample usage for running the graph builder:
+
+`python -m neural_structured_learning.tools.build_graph` [*flags*]
+*embedding_file.tfr... output_graph.tsv*
+
+Sample usage for preparing input for graph-based NSL:
+
+`python -m neural_structured_learning.tools.pack_nbrs` [*flags*]
+*labeled.tfr unlabeled.tfr graph.tsv output.tfr*
+
+For details about these programs' flags, run these commands:
 
 ```sh
-$ python -m neural_structured_learning.tools.build_graph ...
-$ python -m neural_structured_learning.tools.pack_nbrs ...
+$ python -m neural_structured_learning.tools.build_graph --help
+$ python -m neural_structured_learning.tools.pack_nbrs --help
 ```
-
-For details on the command-line usage for these programs, see the
-`nsl.tools.build_graph` and `nsl.tools.pack_nbrs` documentation.
 """
 
 from neural_structured_learning.tools.build_graph import build_graph
+from neural_structured_learning.tools.build_graph import build_graph_from_config
 from neural_structured_learning.tools.graph_utils import add_edge
 from neural_structured_learning.tools.graph_utils import add_undirected_edges
 from neural_structured_learning.tools.graph_utils import read_tsv_graph