build_graph_test.py

# Copyright 2019 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tests for neural_structured_learning.tools.build_graph."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import cmath

from absl.testing import absltest
from neural_structured_learning.tools import build_graph as build_graph_lib
from neural_structured_learning.tools import graph_utils
import six
import tensorflow as tf

from google.protobuf import text_format


# These embeddings in R^3 are chosen so that any pair of them has a cosine
# similaritt of 0.5.
r3_embeddings = [
    """
      features {
        feature {
          key: "id"
          value: { bytes_list { value: [ "A" ] } }
        }
        feature {
          key: "embedding"
          value: { float_list { value: [ 1, 1, 0 ] } }
        }
      }
    """, """
      features {
        feature {
          key: "id"
          value: { bytes_list { value: [ "B" ] } }
        }
        feature {
          key: "embedding"
          value: { float_list { value: [ 1, 0, 1] } }
        }
      }
    """, """
      features {
        feature {
          key: "id"
          value: { bytes_list { value: [ "C" ] } }
        }
        feature {
          key: "embedding"
          value: { float_list { value: [ 0, 1, 1] } }
        }
      }
    """
]


def write_embeddings(embeddings, embedding_output_path):
  """Writes the given 'embeddings' to a TFRecord file at the given path."""
  with tf.io.TFRecordWriter(embedding_output_path) as writer:
    for example_str in embeddings:
      example = text_format.Parse(example_str, tf.train.Example())
      writer.write(example.SerializeToString())


class BuildGraphTest(absltest.TestCase):
  # These embedding vectors are chosen so that the cosine of the angle between
  # each pair of them is 0.5.

  def _create_embedding_file(self):
    return self.create_tempfile('embeddings.tfr').full_path

  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)

  def testBuildGraphInvalidLshRoundsValue(self):
    with self.assertRaises(ValueError):
      build_graph_lib.build_graph([], None, lsh_splits=1, lsh_rounds=0)

  def testBuildGraphNoThresholdingNoLSH(self):
    """All edges whose weight is greater than 0 are retained."""
    embeddings = r3_embeddings
    embedding_path = self._create_embedding_file()
    write_embeddings(embeddings, embedding_path)
    graph_path = self._create_graph_file()
    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, {
            'A': {
                'B': 0.5,
                'C': 0.5
            },
            'B': {
                'A': 0.5,
                'C': 0.5
            },
            'C': {
                'A': 0.5,
                'B': 0.5
            }
        })

  def testBuildGraphWithThresholdingNoLSH(self):
    """Edges below the similarity threshold are not part of the graph."""
    embeddings = r3_embeddings
    embedding_path = self._create_embedding_file()
    write_embeddings(embeddings, embedding_path)
    graph_path = self._create_graph_file()
    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, {})

  def _build_test_embeddings(self, num_points):
    """Construct `num_points` 2D points arranged equiangularly about the origin.

    Use a magnitude multiplier of 1.1 to test that the cosine similarity
    function is ignoring the magnitutes of the vectors (i.e., is normalizing
    the vectors before taking their dot product).

    Args:
      num_points: the number of points to create. The first is `(1.0, 0.0)`.

    Returns:
      A pair containing the resulting embeddings and the cosine similarity
      between adjacent points.
    """
    rotation = cmath.rect(1.1, 2 * cmath.pi / num_points)
    vector = 1.0 + 0.0j
    embeddings = []
    for node_id in range(num_points):
      embedding = tf.train.Example()
      embedding.features.feature['id'].bytes_list.value.append(
          'id_{}'.format(node_id).encode('utf8'))
      values = embedding.features.feature['embedding'].float_list.value
      values.append(vector.real)
      values.append(vector.imag)
      embeddings.append(text_format.MessageToString(embedding))
      vector = vector * rotation
    # Cosine similarity between adjacent points = 0.951057 for 20 points.
    adjacent_similarity = round(rotation.real / abs(rotation), 6)
    return (embeddings, adjacent_similarity)

  def testBuildGraphWithThresholdWithLSHInsufficientLSHRounds(self):
    """Tests that some edges are lost with insufficient LSH rounds."""
    # Construct the embeddings and write them to a file.
    num_points = 20
    (embeddings, _) = self._build_test_embeddings(num_points)
    embedding_path = self._create_embedding_file()
    write_embeddings(embeddings, embedding_path)

    # Build the graph, and read the results into a dictionary.
    graph_path = self._create_graph_file()
    build_graph_lib.build_graph([embedding_path],
                                graph_path,
                                similarity_threshold=0.9,
                                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
    actual_edge_cnt = 0
    for (unused_src_id, tgt_dict) in six.iteritems(g_actual):
      actual_edge_cnt += len(tgt_dict)
    self.assertEqual(actual_edge_cnt, 2 * len(embeddings) - 8,
                     'Expected some edges not to have been found.')

  def testBuildGraphWithThresholdWithLSHSufficientLSHRounds(self):
    """Tests the case where we use (multiple rounds of) LSH bucketing."""
    # Construct the embeddings and write them to a file.
    num_points = 20
    (embeddings, adjacent_similarity) = self._build_test_embeddings(num_points)
    embedding_path = self._create_embedding_file()
    write_embeddings(embeddings, embedding_path)

    # Build the graph, and read the results into a dictionary.
    graph_path = self._create_graph_file()
    build_graph_lib.build_graph([embedding_path],
                                graph_path,
                                similarity_threshold=0.9,
                                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
    # point before it and the point after it in the 'embeddings' sequence.
    # That's because the cosine similarity of adjacent points is ~0.951057,
    # while between every other point it is ~0.809017 (which is below the
    # similarity threshold of 0.9).
    g_expected = {}
    for node_id in range(num_points):
      t_dict = g_expected.setdefault('id_{}'.format(node_id), {})
      t_dict['id_{}'.format((node_id - 1) % num_points)] = adjacent_similarity
      t_dict['id_{}'.format((node_id + 1) % num_points)] = adjacent_similarity
    self.assertDictEqual(g_actual, g_expected)


if __name__ == '__main__':
  # Ensure TF 2.0 behavior even if TF 1.X is installed.
  tf.compat.v1.enable_v2_behavior()
  absltest.main()