//
//  FineTuningJobParameters.swift
//
//
//  Created by James Rochabrun on 10/17/23.
//

import Foundation

/// [Creates a job](https://platform.openai.com/docs/api-reference/fine-tuning/create) that fine-tunes a specified model from a given dataset.
/// Response includes details of the enqueued job including job status and the name of the fine-tuned models once complete.
public struct FineTuningJobParameters: Encodable {
  public init(
    model: Model,
    trainingFile: String,
    hyperparameters: HyperParameters? = nil,
    suffix: String? = nil,
    validationFile: String? = nil,
    integrations: [Integration]? = nil,
    seed: Int? = nil)
  {
    self.model = model.rawValue
    self.trainingFile = trainingFile
    self.hyperparameters = hyperparameters
    self.suffix = suffix
    self.validationFile = validationFile
    self.integrations = integrations
    self.seed = seed
  }

  /// Fine-tuning is [currently available](https://platform.openai.com/docs/guides/fine-tuning/what-models-can-be-fine-tuned) for the following models:
  /// gpt-3.5-turbo-0613 (recommended)
  /// babbage-002
  /// davinci-002
  /// OpenAI expects gpt-3.5-turbo to be the right model for most users in terms of results and ease of use, unless you are migrating a legacy fine-tuned model.
  public enum Model: String {
    case gpt35 = "gpt-3.5-turbo-0613" /// recommended
    case babbage002 = "babbage-002"
    case davinci002 = "davinci-002"
  }

  public struct HyperParameters: Encodable {
    /// The number of epochs to train the model for. An epoch refers to one full cycle through the training dataset.
    /// Defaults to auto.
    let nEpochs: Int?

    public init(
      nEpochs: Int?)
    {
      self.nEpochs = nEpochs
    }

    enum CodingKeys: String, CodingKey {
      case nEpochs = "n_epochs"
    }
  }

  public struct Integration: Encodable {
    /// The type of integration to enable. Currently, only "wandb" (Weights and Biases) is supported.
    let type: String

    let wandb: Wandb

    public struct Wandb: Encodable {
      /// The name of the project that the new run will be created under.
      let project: String
      /// A display name to set for the run. If not set, we will use the Job ID as the name.
      let name: String?
      /// The entity to use for the run. This allows you to set the team or username of the WandB user that you would like associated with the run. If not set, the default entity for the registered WandB API key is used.
      let entity: String?
      /// A list of tags to be attached to the newly created run. These tags are passed through directly to WandB. Some default tags are generated by OpenAI: "openai/finetune", "openai/{base-model}", "openai/{ftjob-abcdef}".
      let tags: [String]?
    }
  }

  enum CodingKeys: String, CodingKey {
    case model
    case trainingFile = "training_file"
    case hyperparameters
    case suffix
    case validationFile = "validation_file"
    case integrations
    case seed
  }

  /// The name of the model to fine-tune. You can select one of the [supported models](https://platform.openai.com/docs/models/overview).
  let model: String
  /// The ID of an uploaded file that contains training data.
  /// See [upload file](https://platform.openai.com/docs/api-reference/files/upload) for how to upload a file.
  /// Your dataset must be formatted as a JSONL file. Additionally, you must upload your file with the purpose fine-tune.
  /// See the [fine-tuning guide](https://platform.openai.com/docs/guides/fine-tuning) for more details.
  let trainingFile: String
  /// The hyperparameters used for the fine-tuning job.
  let hyperparameters: HyperParameters?
  /// A string of up to 18 characters that will be added to your fine-tuned model name.
  /// For example, a suffix of "custom-model-name" would produce a model name like ft:gpt-3.5-turbo:openai:custom-model-name:7p4lURel.
  /// Defaults to null.
  let suffix: String?
  /// The ID of an uploaded file that contains validation data.
  /// If you provide this file, the data is used to generate validation metrics periodically during fine-tuning. These metrics can be viewed in the fine-tuning results file. The same data should not be present in both train and validation files.
  /// Your dataset must be formatted as a JSONL file. You must upload your file with the purpose fine-tune.
  /// See the [fine-tuning guide](https://platform.openai.com/docs/guides/fine-tuning) for more details.
  let validationFile: String?
  /// A list of integrations to enable for your fine-tuning job.
  let integrations: [Integration]?
  /// The seed controls the reproducibility of the job. Passing in the same seed and job parameters should produce the same results, but may differ in rare cases. If a seed is not specified, one will be generated for you.
  let seed: Int?
}
