chromium/third_party/mediapipe/src/mediapipe/tasks/cc/genai/inference/proto/sampler_params.proto

// Copyright 2024 The ODML Authors.
//
// 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
//
//      http://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.

syntax = "proto3";

package odml.infra.proto;

option java_package = "com.google.odml.infra.proto";
option java_outer_classname = "SamplerParametersProto";

// Sampling parameters used to configure the decoding strategy.
message SamplerParameters {
  enum Type {
    TYPE_UNSPECIFIED = 0;
    // Probabilistically pick among the top k tokens.
    TOP_K = 1;
    // Probabilistically pick among the tokens such that the sum is greater
    // than or equal to p tokens after first performing top-k sampling.
    TOP_P = 2;
    // Pick the token with maximum logit (i.e., argmax).
    GREEDY = 3;
  }
  // The type of sampling used to pick the winning token. Ignored on the GPU
  // path, which defaults to combining top-k and top-p sampling.
  Type type = 1;
  // The value of k determines how many of the top-k logits are used during
  // sampling. This field is relevant for TOP_K and TOP_P sampling.
  int32 k = 2;
  // The value of p determines the probability threshold used in TOP_P
  // sampling.
  float p = 3;
  // The temperature used to scale the logits before computing softmax.
  float temperature = 4;
  // The seed used to initialize the random number generator.
  optional int32 seed = 5;
}