Add uniform sampler and combine it with log_uniform sampler into negative_sampler.

PiperOrigin-RevId: 368921933

Author: Neural-Link Team

research/carls/BUILD

@@ -153,7 +153,7 @@ cc_library(
         ":knowledge_bank_service_cc_grpc_proto",
         "//research/carls/candidate_sampling:brute_force_topk_sampler",
         "//research/carls/candidate_sampling:candidate_sampler",
-        "//research/carls/candidate_sampling:log_uniform_sampler",
+        "//research/carls/candidate_sampling:negative_sampler",
         "//research/carls/gradient_descent:gradient_descent_optimizer",
         "//research/carls/knowledge_bank",
         "//research/carls/knowledge_bank:in_proto_knowledge_bank",

research/carls/candidate_sampling/BUILD

@@ -84,8 +84,8 @@ cc_test(
 )
 
 cc_library(
-    name = "log_uniform_sampler",
-    srcs = ["log_uniform_sampler.cc"],
+    name = "negative_sampler",
+    srcs = ["negative_sampler.cc"],
     deps = [
         ":candidate_sampler",
         "//research/carls/base:proto_helper",
@@ -98,8 +98,8 @@ cc_library(
 )
 
 cc_test(
-    name = "log_uniform_sampler_test",
-    srcs = ["log_uniform_sampler_test.cc"],
+    name = "negative_sampler_test",
+    srcs = ["negative_sampler_test.cc"],
     deps = [
         ":candidate_sampler",
         ":log_uniform_sampler",

research/carls/candidate_sampling/candidate_sampler_config.proto

@@ -15,9 +15,18 @@ message CandidateSamplerConfig {
 // A sampler used in the sampled softmax/logistic layer.
 // See https://www.tensorflow.org/extras/candidate_sampling.pdf for more
 // details.
-message LogUniformSamplerConfig {
+message NegativeSamplerConfig {
   // If the sampled key should be unique or not.
   bool unique = 1;
+
+  enum Sampler {
+    UNKNOWN = 0;
+    // Sample uniformly among both positive and negative data.
+    UNIFORM = 1;
+    // Sample log uniformly among both positive and negative data.
+    LOG_UNIFORM = 2;
+  }
+  Sampler sampler = 2;
 }
 
 // Types of similarity between two embeddings.

research/carls/candidate_sampling/candidate_sampler_config_builder.py

@@ -18,8 +18,29 @@
 from research.carls.candidate_sampling import candidate_sampler_config_pb2 as cs_config_pb2
 
 
-def log_uniform_sampler(unique: bool) -> cs_config_pb2.LogUniformSamplerConfig:
-  return cs_config_pb2.LogUniformSamplerConfig(unique=unique)
+def negative_sampler(unique: bool,
+                     algorithm) -> cs_config_pb2.NegativeSamplerConfig:
+  """Builds a NegativeSamplerConfig from given input.
+
+  Args:
+    unique: a bool indicating if the samples should be unique.
+    algorithm: the sampler algorithm defined by NegativeSamplerConfig.Sampler.
+
+  Returns:
+    A NegativeSamplerConfig proto.
+  """
+  if isinstance(algorithm, typing.Text):
+    algorithm = cs_config_pb2.NegativeSamplerConfig.Sampler.Value(algorithm)
+  if isinstance(algorithm, int):
+    if algorithm not in [
+        cs_config_pb2.NegativeSamplerConfig.UNIFORM,
+        cs_config_pb2.NegativeSamplerConfig.LOG_UNIFORM
+    ]:
+      raise ValueError('Invalid sampler type.')
+  else:
+    raise ValueError('Invalid input: %r' % algorithm)
+
+  return cs_config_pb2.NegativeSamplerConfig(unique=unique, sampler=algorithm)
 
 
 def brute_force_topk_sampler(
@@ -53,7 +74,7 @@ def build_candidate_sampler_config(
   """Builds a CandidateSamplerConfig from given sampler.
 
   Args:
-    sampler: an instance of LogUniformSamplerConfig or
+    sampler: an instance of NegativeSamplerConfig or
       BruteForceTopkSamplerConfig.
 
   Returns:
@@ -62,10 +83,10 @@ def build_candidate_sampler_config(
   Raises:
     ValueError if `sampler` is not valid.
   """
-  if not (isinstance(sampler, cs_config_pb2.LogUniformSamplerConfig) or
+  if not (isinstance(sampler, cs_config_pb2.NegativeSamplerConfig) or
           isinstance(sampler, cs_config_pb2.BruteForceTopkSamplerConfig)):
     raise ValueError(
-        'sampler must be one of LogUniformSamplerConfig or BruteForceTopkSamplerConfig'
+        'sampler must be one of NegativeSamplerConfig or BruteForceTopkSamplerConfig'
     )
 
   sampler_config = cs_config_pb2.CandidateSamplerConfig()

research/carls/candidate_sampling/candidate_sampler_config_builder_test.py

@@ -21,13 +21,45 @@
 
 class CandidateSamplerConfigBuilderTest(tf.test.TestCase):
 
-  def test_log_uniform_sampler(self):
-    self.assertProtoEquals("""
+  def test_negative_sampler(self):
+    self.assertProtoEquals(
+        """
       unique: true
-    """, cs_config_builder.log_uniform_sampler(True))
-    self.assertProtoEquals("""
+      sampler: UNIFORM
+    """,
+        cs_config_builder.negative_sampler(
+            True, cs_config_pb2.NegativeSamplerConfig.UNIFORM))
+    self.assertProtoEquals(
+        """
+      unique: false
+      sampler: UNIFORM
+    """,
+        cs_config_builder.negative_sampler(
+            False, cs_config_pb2.NegativeSamplerConfig.UNIFORM))
+    self.assertProtoEquals(
+        """
+      unique: true
+      sampler: LOG_UNIFORM
+    """,
+        cs_config_builder.negative_sampler(
+            True, cs_config_pb2.NegativeSamplerConfig.LOG_UNIFORM))
+    self.assertProtoEquals(
+        """
       unique: false
-    """, cs_config_builder.log_uniform_sampler(False))
+      sampler: LOG_UNIFORM
+    """,
+        cs_config_builder.negative_sampler(
+            False, cs_config_pb2.NegativeSamplerConfig.LOG_UNIFORM))
+    self.assertProtoEquals(
+        """
+      unique: false
+      sampler: UNIFORM
+    """, cs_config_builder.negative_sampler(False, 'UNIFORM'))
+    self.assertProtoEquals(
+        """
+      unique: true
+      sampler: LOG_UNIFORM
+    """, cs_config_builder.negative_sampler(True, 'LOG_UNIFORM'))
 
   def test_brute_force_topk_sampler_success(self):
     self.assertProtoEquals("""
@@ -71,13 +103,14 @@ def test_build_candidate_sampler_config_success(self):
     self.assertProtoEquals(
         """
         extension {
-          [type.googleapis.com/carls.candidate_sampling.LogUniformSamplerConfig] {
+          [type.googleapis.com/carls.candidate_sampling.NegativeSamplerConfig] {
             unique: true
+            sampler: UNIFORM
           }
         }
     """,
         cs_config_builder.build_candidate_sampler_config(
-            cs_config_builder.log_uniform_sampler(True)))
+            cs_config_builder.negative_sampler(True, 'UNIFORM')))
 
   def test_build_candidate_sampler_config_failed(self):
     with self.assertRaises(ValueError):