// Copyright 2024 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
//
//     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 google.ai.generativelanguage.v1beta;

import "google/api/field_behavior.proto";
import "google/api/resource.proto";
import "google/protobuf/timestamp.proto";

option go_package = "cloud.google.com/go/ai/generativelanguage/apiv1beta/generativelanguagepb;generativelanguagepb";
option java_multiple_files = true;
option java_outer_classname = "TunedModelProto";
option java_package = "com.google.ai.generativelanguage.v1beta";

// A fine-tuned model created using ModelService.CreateTunedModel.
message TunedModel {
  option (google.api.resource) = {
    type: "generativelanguage.googleapis.com/TunedModel"
    pattern: "tunedModels/{tuned_model}"
    plural: "tunedModels"
    singular: "tunedModel"
  };

  // The state of the tuned model.
  enum State {
    // The default value. This value is unused.
    STATE_UNSPECIFIED = 0;

    // The model is being created.
    CREATING = 1;

    // The model is ready to be used.
    ACTIVE = 2;

    // The model failed to be created.
    FAILED = 3;
  }

  // The model used as the starting point for tuning.
  oneof source_model {
    // Optional. TunedModel to use as the starting point for training the new
    // model.
    TunedModelSource tuned_model_source = 3
        [(google.api.field_behavior) = OPTIONAL];

    // Immutable. The name of the `Model` to tune.
    // Example: `models/text-bison-001`
    string base_model = 4 [
      (google.api.field_behavior) = IMMUTABLE,
      (google.api.resource_reference) = {
        type: "generativelanguage.googleapis.com/Model"
      }
    ];
  }

  // Output only. The tuned model name. A unique name will be generated on
  // create. Example: `tunedModels/az2mb0bpw6i` If display_name is set on
  // create, the id portion of the name will be set by concatenating the words
  // of the display_name with hyphens and adding a random portion for
  // uniqueness. Example:
  //     display_name = "Sentence Translator"
  //     name = "tunedModels/sentence-translator-u3b7m"
  string name = 1 [(google.api.field_behavior) = OUTPUT_ONLY];

  // Optional. The name to display for this model in user interfaces.
  // The display name must be up to 40 characters including spaces.
  string display_name = 5 [(google.api.field_behavior) = OPTIONAL];

  // Optional. A short description of this model.
  string description = 6 [(google.api.field_behavior) = OPTIONAL];

  // Optional. Controls the randomness of the output.
  //
  // Values can range over `[0.0,1.0]`, inclusive. A value closer to `1.0` will
  // produce responses that are more varied, while a value closer to `0.0` will
  // typically result in less surprising responses from the model.
  //
  // This value specifies default to be the one used by the base model while
  // creating the model.
  optional float temperature = 11 [(google.api.field_behavior) = OPTIONAL];

  // Optional. For Nucleus sampling.
  //
  // Nucleus sampling considers the smallest set of tokens whose probability
  // sum is at least `top_p`.
  //
  // This value specifies default to be the one used by the base model while
  // creating the model.
  optional float top_p = 12 [(google.api.field_behavior) = OPTIONAL];

  // Optional. For Top-k sampling.
  //
  // Top-k sampling considers the set of `top_k` most probable tokens.
  // This value specifies default to be used by the backend while making the
  // call to the model.
  //
  // This value specifies default to be the one used by the base model while
  // creating the model.
  optional int32 top_k = 13 [(google.api.field_behavior) = OPTIONAL];

  // Output only. The state of the tuned model.
  State state = 7 [(google.api.field_behavior) = OUTPUT_ONLY];

  // Output only. The timestamp when this model was created.
  google.protobuf.Timestamp create_time = 8
      [(google.api.field_behavior) = OUTPUT_ONLY];

  // Output only. The timestamp when this model was updated.
  google.protobuf.Timestamp update_time = 9
      [(google.api.field_behavior) = OUTPUT_ONLY];

  // Required. The tuning task that creates the tuned model.
  TuningTask tuning_task = 10 [(google.api.field_behavior) = REQUIRED];
}

// Tuned model as a source for training a new model.
message TunedModelSource {
  // Immutable. The name of the `TunedModel` to use as the starting point for
  // training the new model.
  // Example: `tunedModels/my-tuned-model`
  string tuned_model = 1 [
    (google.api.field_behavior) = IMMUTABLE,
    (google.api.resource_reference) = {
      type: "generativelanguage.googleapis.com/TunedModel"
    }
  ];

  // Output only. The name of the base `Model` this `TunedModel` was tuned from.
  // Example: `models/text-bison-001`
  string base_model = 2 [
    (google.api.field_behavior) = OUTPUT_ONLY,
    (google.api.resource_reference) = {
      type: "generativelanguage.googleapis.com/Model"
    }
  ];
}

// Tuning tasks that create tuned models.
message TuningTask {
  // Output only. The timestamp when tuning this model started.
  google.protobuf.Timestamp start_time = 1
      [(google.api.field_behavior) = OUTPUT_ONLY];

  // Output only. The timestamp when tuning this model completed.
  google.protobuf.Timestamp complete_time = 2
      [(google.api.field_behavior) = OUTPUT_ONLY];

  // Output only. Metrics collected during tuning.
  repeated TuningSnapshot snapshots = 3
      [(google.api.field_behavior) = OUTPUT_ONLY];

  // Required. Input only. Immutable. The model training data.
  Dataset training_data = 4 [
    (google.api.field_behavior) = INPUT_ONLY,
    (google.api.field_behavior) = REQUIRED,
    (google.api.field_behavior) = IMMUTABLE
  ];

  // Immutable. Hyperparameters controlling the tuning process. If not provided,
  // default values will be used.
  Hyperparameters hyperparameters = 5 [(google.api.field_behavior) = IMMUTABLE];
}

// Hyperparameters controlling the tuning process. Read more at
// https://ai.google.dev/docs/model_tuning_guidance
message Hyperparameters {
  // Options for specifying learning rate during tuning.
  oneof learning_rate_option {
    // Optional. Immutable. The learning rate hyperparameter for tuning.
    // If not set, a default of 0.001 or 0.0002 will be calculated based on the
    // number of training examples.
    float learning_rate = 16 [
      (google.api.field_behavior) = IMMUTABLE,
      (google.api.field_behavior) = OPTIONAL
    ];

    // Optional. Immutable. The learning rate multiplier is used to calculate a
    // final learning_rate based on the default (recommended) value. Actual
    // learning rate := learning_rate_multiplier * default learning rate Default
    // learning rate is dependent on base model and dataset size. If not set, a
    // default of 1.0 will be used.
    float learning_rate_multiplier = 17 [
      (google.api.field_behavior) = IMMUTABLE,
      (google.api.field_behavior) = OPTIONAL
    ];
  }

  // Immutable. The number of training epochs. An epoch is one pass through the
  // training data. If not set, a default of 5 will be used.
  optional int32 epoch_count = 14 [(google.api.field_behavior) = IMMUTABLE];

  // Immutable. The batch size hyperparameter for tuning.
  // If not set, a default of 4 or 16 will be used based on the number of
  // training examples.
  optional int32 batch_size = 15 [(google.api.field_behavior) = IMMUTABLE];
}

// Dataset for training or validation.
message Dataset {
  // Inline data or a reference to the data.
  oneof dataset {
    // Optional. Inline examples.
    TuningExamples examples = 1 [(google.api.field_behavior) = OPTIONAL];
  }
}

// A set of tuning examples. Can be training or validation data.
message TuningExamples {
  // Required. The examples. Example input can be for text or discuss, but all
  // examples in a set must be of the same type.
  repeated TuningExample examples = 1 [(google.api.field_behavior) = REQUIRED];
}

// A single example for tuning.
message TuningExample {
  // The input to the model for this example.
  oneof model_input {
    // Optional. Text model input.
    string text_input = 1 [(google.api.field_behavior) = OPTIONAL];
  }

  // Required. The expected model output.
  string output = 3 [(google.api.field_behavior) = REQUIRED];
}

// Record for a single tuning step.
message TuningSnapshot {
  // Output only. The tuning step.
  int32 step = 1 [(google.api.field_behavior) = OUTPUT_ONLY];

  // Output only. The epoch this step was part of.
  int32 epoch = 2 [(google.api.field_behavior) = OUTPUT_ONLY];

  // Output only. The mean loss of the training examples for this step.
  float mean_loss = 3 [(google.api.field_behavior) = OUTPUT_ONLY];

  // Output only. The timestamp when this metric was computed.
  google.protobuf.Timestamp compute_time = 4
      [(google.api.field_behavior) = OUTPUT_ONLY];
}