knowledge_bank_service.proto

syntax = "proto3";

package carls;

import "research/carls/candidate_sampling/candidate_sampler_config.proto";
import "research/carls/dynamic_embedding_config.proto";
import "research/carls/embedding.proto";
import "research/carls/memory_store/memory_store_config.proto";

message StartSessionRequest {
  // The config for processing embedding lookup/update.
  DynamicEmbeddingConfig config = 1;

  // The name of the Op that uses dynamic embedding.
  // A name and a config should uniquely identify a session.
  string name = 2;
}

message StartSessionResponse {
  // A unique key for DES to identify this session.
  bytes session_handle = 1;
}

message LookupRequest {
  // A handle to identify which session to use.
  bytes session_handle = 1;

  // List of lookup keys.
  repeated string key = 2;

  // If update = true, Lookup should allocate a new entry for a new key and
  // also count frequency, often used in feedforward path in training.
  // Otherwise, it should just be a lookup without changing any internal
  // information, often used in inference.
  bool update = 3;
}

message LookupResponse {
  // Maps from keys to their embedding.
  map<string, EmbeddingVectorProto> embedding_table = 1;
}

message UpdateRequest {
  // A handle to identify which session to use.
  bytes session_handle = 1;

  // A batch of keys and their new values to be updated.
  map<string, EmbeddingVectorProto> values = 2;

  // A batch of keys and gradients to be updated.
  map<string, EmbeddingVectorProto> gradients = 3;
}

message UpdateResponse {}

message SampleRequest {
  // A handle to identify which session to use.
  bytes session_handle = 1;

  // A batch of sample contexts.
  repeated candidate_sampling.SampleContext sample_context = 2;

  // Number of returned samples per sample context.
  int32 num_samples = 3;

  // If true, allocate new embeddings for positive keys that are not in the
  // knowledge bank.
  bool update = 4;
}

message SampleResponse {
  message Samples {
    // Sampled results.
    repeated candidate_sampling.SampledResult sampled_result = 1;
  }

  // A batch of results corresponding to the input sample context.
  repeated Samples samples = 1;
}

message MemoryLookupRequest {
  // A handle to identify which session to use.
  bytes session_handle = 1;

  // List of lookup activations.
  repeated EmbeddingVectorProto input = 2;

  enum LookupMode {
    DEFAULT_UNKNOWN = 0;
    // Lookup without changing existing memory data.
    LOOKUP_WITHOUT_UPDATE = 1;
    // Lookup keys are used to update existing memory data, but new clusters
    // would not be created.
    LOOKUP_WITH_UPDATE = 2;
    // Lookup keys are used to update existing memory data, and a new cluster
    // can be created if the lookup input is too far away from existing ones.
    LOOKUP_WITH_GROW = 3;
  }
  LookupMode mode = 3;
}

message MemoryLookupResponse {
  // Memory lookup result for each input.
  repeated memory_store.MemoryLookupResult memory_lookup_result = 1;
}

message ExportRequest {
  // A handle to identify which session to use.
  bytes session_handle = 1;

  // Path to the export directory.
  string export_directory = 2;
}

message ExportResponse {
  oneof saved_path {
    // The saved path for knowledge bank.
    string knowledge_bank_saved_path = 1;

    // The saved path for memory store. Used by DynamicMemory.
    string memory_store_saved_path = 2;
  }
}

message ImportRequest {
  // A handle to identify which session to use.
  bytes session_handle = 1;

  oneof saved_path {
    // The saved path for knowledge bank.
    string knowledge_bank_saved_path = 2;

    // The saved path for memory store. Used by DynamicMemory.
    string memory_store_saved_path = 3;
  }
}

message ImportResponse {}

// KnowledgeBankService defines the service for handling embedding lookup,
// updates and samples.
service KnowledgeBankService {
  // Starts a session with the DES server.
  rpc StartSession(StartSessionRequest) returns (StartSessionResponse);

  // Lookups the embeddings from a given batch of keys.
  rpc Lookup(LookupRequest) returns (LookupResponse);

  // Updates the embedding value for a given batch of keys.
  rpc Update(UpdateRequest) returns (UpdateResponse);

  // Samples the embeddings value from given context.
  rpc Sample(SampleRequest) returns (SampleResponse);

  // Looks up a dynamic memory layer and returns the corresponding memory data.
  rpc MemoryLookup(MemoryLookupRequest) returns (MemoryLookupResponse);

  // Exports current model to a given directory with timestamped subdir.
  rpc Export(ExportRequest) returns (ExportResponse);

  // Imports the state of DES for a given session_handle.
  rpc Import(ImportRequest) returns (ImportResponse);
}