Enable Candidate Sampling Part II: add Sample RPC

PiperOrigin-RevId: 366610296

Author: Neural-Link Team

research/carls/BUILD

@@ -66,6 +66,7 @@ carls_cc_proto_library(
     name = "dynamic_embedding_config_cc_proto",
     srcs = ["dynamic_embedding_config.proto"],
     deps = [
+        "//research/carls/candidate_sampling:candidate_sampler_config_cc_proto",
         "//research/carls/gradient_descent:gradient_descent_config_cc_proto",
         "//research/carls/knowledge_bank:knowledge_bank_config_cc_proto",
     ],
@@ -76,6 +77,7 @@ carls_py_proto_library(
     srcs = ["dynamic_embedding_config.proto"],
     deps = [
         ":dynamic_embedding_config_cc_proto",
+        "//research/carls/candidate_sampling:candidate_sampler_config_py_pb2",
         "//research/carls/gradient_descent:gradient_descent_config_py_pb2",
         "//research/carls/knowledge_bank:knowledge_bank_config_py_pb2",
     ],
@@ -87,6 +89,7 @@ carls_cc_proto_library(
     deps = [
         ":dynamic_embedding_config_cc_proto",
         ":embedding_cc_proto",
+        "//research/carls/candidate_sampling:candidate_sampler_config_cc_proto",
     ],
 )
 
@@ -148,6 +151,9 @@ cc_library(
     hdrs = ["knowledge_bank_grpc_service.h"],
     deps = [
         ":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/gradient_descent:gradient_descent_optimizer",
         "//research/carls/knowledge_bank",
         "//research/carls/knowledge_bank:in_proto_knowledge_bank",

research/carls/dynamic_embedding_config.proto

@@ -2,6 +2,7 @@ syntax = "proto3";
 
 package carls;
 
+import "research/carls/candidate_sampling/candidate_sampler_config.proto";
 import "research/carls/gradient_descent/gradient_descent_config.proto";
 import "research/carls/knowledge_bank/knowledge_bank_config.proto";
 
@@ -14,4 +15,6 @@ message DynamicEmbeddingConfig {
   KnowledgeBankConfig knowledge_bank_config = 2;
 
   GradientDescentConfig gradient_descent_config = 3;
+
+  candidate_sampling.CandidateSamplerConfig candidate_sampler_config = 4;
 }

research/carls/knowledge_bank_grpc_service.cc

@@ -176,6 +176,37 @@ Status KnowledgeBankGrpcServiceImpl::Update(grpc::ServerContext* context,
   return Status::OK;
 }
 
+grpc::Status KnowledgeBankGrpcServiceImpl::Sample(grpc::ServerContext* context,
+                                                  const SampleRequest* request,
+                                                  SampleResponse* response) {
+  if (request->session_handle().empty()) {
+    return Status(StatusCode::INVALID_ARGUMENT, "session_handle is empty.");
+  }
+  if (request->sample_context().empty()) {
+    return Status(StatusCode::INVALID_ARGUMENT, "No sample context.");
+  }
+  const auto status = StartSessionIfNecessary(request->session_handle());
+  if (!status.ok()) {
+    return status;
+  }
+  absl::MutexLock lock(&map_mu_);
+  const auto& knowledge_bank = *kb_map_[request->session_handle()];
+
+  for (const auto& sample_context : request->sample_context()) {
+    std::vector<candidate_sampling::SampledResult> results;
+    auto status = cs_map_[request->session_handle()]->Sample(
+        knowledge_bank, sample_context, request->num_samples(), &results);
+    if (!status.ok()) {
+      return ToGrpcStatus(status);
+    }
+    auto* samples = response->add_samples();
+    for (auto& result : results) {
+      *samples->add_sampled_result() = std::move(result);
+    }
+  }
+  return Status::OK;
+}
+
 Status KnowledgeBankGrpcServiceImpl::Export(grpc::ServerContext* context,
                                             const ExportRequest* request,
                                             ExportResponse* response) {
@@ -247,6 +278,15 @@ Status KnowledgeBankGrpcServiceImpl::StartSessionIfNecessary(
     }
     gd_map_[session_handle] = std::move(optimizer);
   }
+  if (request.config().has_candidate_sampler_config() &&
+      !cs_map_.contains(session_handle)) {
+    auto sampler = candidate_sampling::SamplerFactory::Make(
+        request.config().candidate_sampler_config());
+    if (sampler == nullptr) {
+      return Status(StatusCode::INTERNAL, "Creating CandidateSampler failed.");
+    }
+    cs_map_[session_handle] = std::move(sampler);
+  }
   return Status::OK;
 }
 

research/carls/knowledge_bank_grpc_service.h

@@ -20,6 +20,7 @@ limitations under the License.
 
 #include "grpcpp/support/status.h"  // net
 #include "absl/synchronization/mutex.h"
+#include "research/carls/candidate_sampling/candidate_sampler.h"
 #include "research/carls/gradient_descent/gradient_descent_optimizer.h"
 #include "research/carls/knowledge_bank/knowledge_bank.h"
 #include "research/carls/knowledge_bank_service.grpc.pb.h"
@@ -53,6 +54,11 @@ class KnowledgeBankGrpcServiceImpl final
                       const UpdateRequest* request,
                       UpdateResponse* response) override;
 
+  // Implements the Sample method of KnowledgeBankService.
+  grpc::Status Sample(grpc::ServerContext* context,
+                      const SampleRequest* request,
+                      SampleResponse* response) override;
+
   // Implements the Export method of KnowledgeBankService.
   grpc::Status Export(grpc::ServerContext* context,
                       const ExportRequest* request,
@@ -77,6 +83,10 @@ class KnowledgeBankGrpcServiceImpl final
   // Maps from session_handle to GradientDescentOptimizer.
   absl::node_hash_map<std::string, std::unique_ptr<GradientDescentOptimizer>>
       gd_map_;
+  // Maps from session_handle to CandidateSampler.
+  absl::node_hash_map<std::string,
+                      std::unique_ptr<candidate_sampling::CandidateSampler>>
+      cs_map_;
 };
 
 }  // namespace carls