Fixes the NSL graph builder to ignore lsh_rounds when lsh_splits < 1.

This fixes a regression introduced in NSL v1.3.0 in which twice the
work was being performed in the default case. As a workaround, just
specify lsh_rounds=1 when lsh_splits=0.

PiperOrigin-RevId: 327063409

Author: Neural-Link Team

neural_structured_learning/tools/build_graph.py

@@ -219,7 +219,10 @@ def _generate_edges(self, embeddings):
       A tuple (source, target, weight) denoting a (directed) edge from 'source'
       to 'target' with the given 'weight'.
     """
-    for lsh_round in range(max(1, self.config.lsh_rounds)):
+    # If lsh_splits < 1, we ignore lsh_rounds and always perform 1 round, since
+    # performing multiple rounds in the case of no splits does not help.
+    rounds = self.config.lsh_rounds if self.config.lsh_splits > 0 else 1
+    for lsh_round in range(rounds):
       start_time = time.time()
       edge_cnt = 0
       bucket_map = self._generate_lsh_buckets(embeddings)

neural_structured_learning/tools/build_graph_test.py

@@ -86,6 +86,10 @@ def _create_embedding_file(self):
   def _create_graph_file(self):
     return self.create_tempfile('graph.tsv').full_path
 
+  def _num_file_lines(self, graph_path):
+    with open(graph_path, 'rU') as f:
+      return sum(1 for _ in f)
+
   def testBuildGraphInvalidLshBitsValue(self):
     with self.assertRaises(ValueError):
       build_graph_lib.build_graph([], None, lsh_splits=-1)
@@ -103,6 +107,7 @@ def testBuildGraphNoThresholdingNoLSH(self):
     build_graph_lib.build_graph([embedding_path],
                                 graph_path,
                                 similarity_threshold=0)
+    self.assertEqual(self._num_file_lines(graph_path), 6)
     g_actual = graph_utils.read_tsv_graph(graph_path)
     self.assertDictEqual(
         g_actual, {
@@ -129,6 +134,7 @@ def testBuildGraphWithThresholdingNoLSH(self):
     build_graph_lib.build_graph([embedding_path],
                                 graph_path,
                                 similarity_threshold=0.51)
+    self.assertEqual(self._num_file_lines(graph_path), 0)
     g_actual = graph_utils.read_tsv_graph(graph_path)
     self.assertDictEqual(g_actual, {})
 
@@ -178,6 +184,7 @@ def testBuildGraphWithThresholdWithLSHInsufficientLSHRounds(self):
                                 lsh_splits=2,
                                 lsh_rounds=1,
                                 random_seed=12345)
+    self.assertEqual(self._num_file_lines(graph_path), num_points * 2 - 8)
     g_actual = graph_utils.read_tsv_graph(graph_path)
 
     # Check that the graph contains fewer than 2 * N edges
@@ -203,6 +210,7 @@ def testBuildGraphWithThresholdWithLSHSufficientLSHRounds(self):
                                 lsh_splits=2,
                                 lsh_rounds=4,
                                 random_seed=12345)
+    self.assertEqual(self._num_file_lines(graph_path), num_points * 2)
     g_actual = graph_utils.read_tsv_graph(graph_path)
 
     # Constuct the expected graph: each point should be a neighbor of the