tinkex

This file is a merged representation of the entire codebase, combined into a single document by Repomix.
The content has been processed where content has been compressed (code blocks are separated by ⋮---- delimiter).

<file_summary>
This section contains a summary of this file.

<purpose>
This file contains a packed representation of the entire repository's contents.
It is designed to be easily consumable by AI systems for analysis, code review,
or other automated processes.
</purpose>

<file_format>
The content is organized as follows:
1. This summary section
2. Repository information
3. Directory structure
4. Repository files (if enabled)
5. Multiple file entries, each consisting of:
  - File path as an attribute
  - Full contents of the file
</file_format>

<usage_guidelines>
- This file should be treated as read-only. Any changes should be made to the
  original repository files, not this packed version.
- When processing this file, use the file path to distinguish
  between different files in the repository.
- Be aware that this file may contain sensitive information. Handle it with
  the same level of security as you would the original repository.
</usage_guidelines>

<notes>
- Some files may have been excluded based on .gitignore rules and Repomix's configuration
- Binary files are not included in this packed representation. Please refer to the Repository Structure section for a complete list of file paths, including binary files
- Files matching patterns in .gitignore are excluded
- Files matching default ignore patterns are excluded
- Content has been compressed - code blocks are separated by ⋮---- delimiter
- Files are sorted by Git change count (files with more changes are at the bottom)
</notes>

</file_summary>

<directory_structure>
async_client_creation.exs
checkpoint_download.exs
checkpoints_management.exs
cli_run_prompt_file.exs
cli_run_text.exs
sampling_basic.exs
sessions_management.exs
training_loop.exs
</directory_structure>

<files>
This section contains the contents of the repository's files.

<file path="async_client_creation.exs">
defmodule Tinkex.Examples.AsyncClientCreation do
  @moduledoc """
  Example demonstrating async client creation patterns.

  Shows how to:
  - Create multiple sampling clients concurrently
  - Use Task.await_many for parallel operations
  """

  alias Tinkex.{ServiceClient, SamplingClient, Config}

  def run do
    IO.puts("=== Tinkex Async Client Creation Example ===\n")

    {:ok, _} = Application.ensure_all_started(:tinkex)

    base_url =
      System.get_env("TINKER_BASE_URL") ||
        Application.get_env(
          :tinkex,
          :base_url,
          "https://tinker.thinkingmachines.dev/services/tinker-prod"
        )

    config =
      Config.new(
        api_key: System.get_env("TINKER_API_KEY") || raise("TINKER_API_KEY required"),
        base_url: base_url
      )

    # Get checkpoint paths from environment or use single base model
    checkpoint_paths =
      case System.get_env("TINKER_CHECKPOINT_PATHS") do
        nil -> []
        paths -> String.split(paths, ",")
      end

    {:ok, service_pid} = ServiceClient.start_link(config: config)

    if length(checkpoint_paths) > 0 do
      create_multiple_clients(service_pid, checkpoint_paths)
    else
      create_single_client_async(service_pid)
    end

    GenServer.stop(service_pid)
    IO.puts("\n=== Example Complete ===")
  end

  defp create_single_client_async(service_pid) do
    IO.puts("Creating sampling client asynchronously...")

    base_model = System.get_env("TINKER_BASE_MODEL") || "meta-llama/Llama-3.2-1B"

    task = ServiceClient.create_sampling_client_async(service_pid, base_model: base_model)

    IO.puts("Task created, awaiting result...")

    case Task.await(task, 30_000) do
      {:ok, pid} ->
        IO.puts("✓ Sampling client created: #{inspect(pid)}")
        GenServer.stop(pid)

      {:error, reason} ->
        IO.puts("✗ Failed: #{inspect(reason)}")
    end
  end

  defp create_multiple_clients(service_pid, checkpoint_paths) do
    IO.puts("Creating #{length(checkpoint_paths)} sampling clients concurrently...\n")

    # Start timing
    start_time = System.monotonic_time(:millisecond)

    # Create tasks for all clients
    tasks =
      Enum.map(checkpoint_paths, fn path ->
        IO.puts("  Starting task for: #{path}")
        SamplingClient.create_async(service_pid, model_path: path)
      end)

    # Wait for all to complete
    IO.puts("\nAwaiting all tasks...")
    results = Task.await_many(tasks, 60_000)

    # Calculate time
    elapsed = System.monotonic_time(:millisecond) - start_time

    # Report results
    IO.puts("\nResults (#{elapsed}ms total):")

    Enum.zip(checkpoint_paths, results)
    |> Enum.each(fn {path, result} ->
      case result do
        {:ok, pid} ->
          IO.puts("  ✓ #{path} -> #{inspect(pid)}")
          GenServer.stop(pid)

        {:error, reason} ->
          IO.puts("  ✗ #{path} -> #{inspect(reason)}")
      end
    end)

    successes =
      Enum.count(results, fn
        {:ok, _} -> true
        _ -> false
      end)

    IO.puts("\nSuccess: #{successes}/#{length(results)}")
  end
end

Tinkex.Examples.AsyncClientCreation.run()
</file>

<file path="checkpoint_download.exs">
defmodule Tinkex.Examples.CheckpointDownloadExample do
  @moduledoc """
  Example demonstrating checkpoint download functionality.

  Shows how to:
  - Download and extract checkpoint archive
  - Track progress during download
  """

  alias Tinkex.{ServiceClient, RestClient, CheckpointDownload, Config, Error}

  def run do
    IO.puts("=== Tinkex Checkpoint Download Example ===\n")

    {:ok, _} = Application.ensure_all_started(:tinkex)

    base_url =
      System.get_env("TINKER_BASE_URL") ||
        Application.get_env(
          :tinkex,
          :base_url,
          "https://tinker.thinkingmachines.dev/services/tinker-prod"
        )

    config =
      Config.new(
        api_key: System.get_env("TINKER_API_KEY") || raise("TINKER_API_KEY required"),
        base_url: base_url
      )

    {:ok, service_pid} = ServiceClient.start_link(config: config)
    {:ok, rest_client} = ServiceClient.create_rest_client(service_pid)

    {checkpoint_path, path_source} =
      case resolve_checkpoint_path(rest_client, service_pid) do
        {:ok, path, source} ->
          {path, source}

        {:error, :no_checkpoints} ->
          raise """
          No checkpoints found for this account. Set TINKER_CHECKPOINT_PATH (e.g., tinker://run-123/weights/0001)
          or create a checkpoint via training before running the example.
          """

        {:error, :no_available_checkpoints} ->
          raise """
          Could not find a downloadable checkpoint automatically. The most recent entries returned 404.
          Set TINKER_CHECKPOINT_PATH to a known-good checkpoint path and rerun the example.
          """

        {:error, {:api_error, %Error{} = error}} ->
          raise """
          Failed to discover checkpoints automatically: #{Error.format(error)}.
          Set TINKER_CHECKPOINT_PATH manually or rerun later.
          """

        {:error, {:api_error, other}} ->
          raise """
          Failed to discover checkpoints automatically: #{inspect(other)}.
          Set TINKER_CHECKPOINT_PATH manually or rerun later.
          """
      end

    if path_source == :auto do
      IO.puts(
        "TINKER_CHECKPOINT_PATH not provided; downloading first available checkpoint:\n  #{checkpoint_path}\n"
      )
    end

    output_dir = System.get_env("TINKER_OUTPUT_DIR") || File.cwd!()

    IO.puts("Downloading checkpoint: #{checkpoint_path}")
    IO.puts("Output directory: #{output_dir}\n")

    # Progress callback
    progress_fn = fn downloaded, total ->
      percent = if total > 0, do: Float.round(downloaded / total * 100, 1), else: 0
      IO.write("\rProgress: #{percent}% (#{format_size(downloaded)} / #{format_size(total)})")
    end

    force? = System.get_env("FORCE") == "true"

    case CheckpointDownload.download(rest_client, checkpoint_path,
           output_dir: output_dir,
           force: force?,
           progress: progress_fn
         ) do
      {:ok, result} ->
        IO.puts("\n\nDownload complete!")
        IO.puts("Extracted to: #{result.destination}")

        if File.exists?(result.destination) do
          files = File.ls!(result.destination)
          IO.puts("\nExtracted files (#{length(files)}):")

          Enum.each(files, fn file ->
            path = Path.join(result.destination, file)
            stat = File.stat!(path)
            IO.puts("  • #{file} (#{format_size(stat.size)})")
          end)
        end

      {:error, {:exists, path}} ->
        IO.puts("\nError: Directory already exists: #{path}")
        IO.puts("Use FORCE=true to overwrite")

      {:error, %Tinkex.Error{status: 404} = error} ->
        IO.write("""

        Error: #{Tinkex.Error.format(error)}
        The checkpoint might no longer exist, or the service returned 404 for the latest entry.
        Try setting TINKER_CHECKPOINT_PATH explicitly to a known-good value or set FORCE=true if the
        directory already exists.
        """)

      {:error, error} ->
        IO.puts("\nError: #{inspect(error)}")
    end

    GenServer.stop(service_pid)
    IO.puts("\n=== Example Complete ===")
  end

  defp format_size(bytes) when bytes < 1024, do: "#{bytes} B"
  defp format_size(bytes) when bytes < 1024 * 1024, do: "#{Float.round(bytes / 1024, 1)} KB"

  defp format_size(bytes) when bytes < 1024 * 1024 * 1024,
    do: "#{Float.round(bytes / 1024 / 1024, 1)} MB"

  defp format_size(bytes), do: "#{Float.round(bytes / 1024 / 1024 / 1024, 2)} GB"

  defp resolve_checkpoint_path(rest_client, service_pid) do
    case System.get_env("TINKER_CHECKPOINT_PATH") do
      nil ->
        case find_available_checkpoint(rest_client) do
          {:ok, path} ->
            {:ok, path, :auto}

          {:error, :no_available_checkpoints} ->
            with {:ok, path} <- create_temporary_checkpoint(service_pid) do
              {:ok, path, :generated}
            end

          {:error, other} ->
            {:error, other}
        end

      path ->
        {:ok, path, :env}
    end
  end

  defp find_available_checkpoint(rest_client, offset \\ 0) do
    case RestClient.list_user_checkpoints(rest_client, limit: 10, offset: offset) do
      {:ok, %Tinkex.Types.CheckpointsListResponse{checkpoints: []}} ->
        if offset == 0, do: {:error, :no_checkpoints}, else: {:error, :no_available_checkpoints}

      {:ok, %Tinkex.Types.CheckpointsListResponse{checkpoints: checkpoints}} ->
        case Enum.reduce_while(checkpoints, {:error, :no_available_checkpoints}, fn ckpt, _acc ->
               case RestClient.get_checkpoint_archive_url(rest_client, ckpt.tinker_path) do
                 {:ok, _} ->
                   {:halt, {:ok, ckpt.tinker_path}}

                 {:error, %Error{status: 404}} ->
                   {:cont, {:error, :no_available_checkpoints}}

                 {:error, other} ->
                   {:halt, {:error, {:api_error, other}}}
               end
             end) do
          {:ok, path} ->
            {:ok, path}

          {:error, :no_available_checkpoints} ->
            find_available_checkpoint(rest_client, offset + length(checkpoints))

          {:error, other} ->
            {:error, other}
        end

      {:error, error} ->
        {:error, {:api_error, error}}
    end
  end

  defp create_temporary_checkpoint(service_pid) do
    base_model = System.get_env("TINKER_BASE_MODEL") || "meta-llama/Llama-3.1-8B"

    with {:ok, training_pid} <-
           ServiceClient.create_lora_training_client(service_pid, base_model: base_model),
         {:ok, save_task} <- TrainingClient.save_weights_for_sampler(training_pid) do
      result = Task.await(save_task, 60_000)
      GenServer.stop(training_pid)

      case result do
        {:ok, %{"path" => path}} when is_binary(path) ->
          {:ok, path}

        {:ok, %{path: path}} when is_binary(path) ->
          {:ok, path}

        {:error, _} = error ->
          error

        other ->
          {:error,
           Error.new(:request_failed, "Unexpected save_weights_for_sampler response",
             data: other
           )}
      end
    else
      {:error, _} = error ->
        error
    end
  end
end

Tinkex.Examples.CheckpointDownloadExample.run()
</file>

<file path="checkpoints_management.exs">
defmodule Tinkex.Examples.CheckpointsManagement do
  @moduledoc """
  Example demonstrating checkpoint management APIs.

  Shows how to:
  - List all user checkpoints
  - List checkpoints for a specific run
  - View checkpoint details
  """

  alias Tinkex.{ServiceClient, RestClient, Config}

  def run do
    IO.puts("=== Tinkex Checkpoint Management Example ===\n")

    {:ok, _} = Application.ensure_all_started(:tinkex)

    base_url =
      System.get_env("TINKER_BASE_URL") ||
        Application.get_env(
          :tinkex,
          :base_url,
          "https://tinker.thinkingmachines.dev/services/tinker-prod"
        )

    config =
      Config.new(
        api_key: System.get_env("TINKER_API_KEY") || raise("TINKER_API_KEY required"),
        base_url: base_url
      )

    {:ok, service_pid} = ServiceClient.start_link(config: config)
    {:ok, rest_client} = ServiceClient.create_rest_client(service_pid)

    # List all user checkpoints
    list_all_checkpoints(rest_client)

    # If run_id is provided, list checkpoints for that run
    if run_id = System.get_env("TINKER_RUN_ID") do
      list_run_checkpoints(rest_client, run_id)
    end

    GenServer.stop(service_pid)
    IO.puts("\n=== Example Complete ===")
  end

  defp list_all_checkpoints(rest_client) do
    IO.puts("--- All User Checkpoints ---")

    case RestClient.list_user_checkpoints(rest_client, limit: 20) do
      {:ok, response} ->
        total =
          if response.cursor,
            do: response.cursor["total_count"],
            else: length(response.checkpoints)

        IO.puts("Found #{length(response.checkpoints)} of #{total || "?"} checkpoints:\n")

        Enum.each(response.checkpoints, fn ckpt ->
          size = if ckpt.size_bytes, do: format_size(ckpt.size_bytes), else: "N/A"
          IO.puts("  #{ckpt.checkpoint_id}")
          IO.puts("    Path: #{ckpt.tinker_path}")
          IO.puts("    Type: #{ckpt.checkpoint_type}")
          IO.puts("    Size: #{size}")
          IO.puts("    Public: #{ckpt.public}")
          IO.puts("    Created: #{ckpt.time}")
          IO.puts("")
        end)

      {:error, error} ->
        IO.puts("Error: #{inspect(error)}")
    end
  end

  defp list_run_checkpoints(rest_client, run_id) do
    IO.puts("\n--- Checkpoints for Run: #{run_id} ---")

    case RestClient.list_checkpoints(rest_client, run_id) do
      {:ok, response} ->
        IO.puts("Found #{length(response.checkpoints)} checkpoints:")

        Enum.each(response.checkpoints, fn ckpt ->
          IO.puts("  • #{ckpt.tinker_path}")
        end)

      {:error, error} ->
        IO.puts("Error: #{inspect(error)}")
    end
  end

  defp format_size(bytes) when bytes < 1024, do: "#{bytes} B"
  defp format_size(bytes) when bytes < 1024 * 1024, do: "#{Float.round(bytes / 1024, 1)} KB"

  defp format_size(bytes) when bytes < 1024 * 1024 * 1024,
    do: "#{Float.round(bytes / 1024 / 1024, 1)} MB"

  defp format_size(bytes), do: "#{Float.round(bytes / 1024 / 1024 / 1024, 2)} GB"
end

Tinkex.Examples.CheckpointsManagement.run()
</file>

<file path="cli_run_prompt_file.exs">
defmodule Tinkex.Examples.CLIRunPromptFile do
  @moduledoc false
  alias Tinkex.CLI

  def run do
    {:ok, _} = Application.ensure_all_started(:tinkex)

    api_key =
      System.get_env("TINKER_API_KEY") ||
        raise "Set TINKER_API_KEY to run this example"

    base_model = System.get_env("TINKER_BASE_MODEL", "meta-llama/Llama-3.1-8B")
    base_url = System.get_env("TINKER_BASE_URL")

    prompt_content =
      System.get_env("TINKER_PROMPT_TOKENS") ||
        System.get_env("TINKER_PROMPT", "Hello from a prompt file")

    tmp_dir = System.tmp_dir!()
    prompt_path = Path.join(tmp_dir, "tinkex_prompt_#{System.unique_integer([:positive])}.txt")
    output_path = Path.join(tmp_dir, "tinkex_output_#{System.unique_integer([:positive])}.json")

    File.write!(prompt_path, prompt_content)

    args =
      [
        "run",
        "--base-model",
        base_model,
        "--prompt-file",
        prompt_path,
        "--json",
        "--output",
        output_path,
        "--api-key",
        api_key
      ]
      |> maybe_add_base_url(base_url)

    IO.puts("Running CLI with prompt file #{prompt_path}")

    case CLI.run(args) do
      {:ok, %{response: _resp}} ->
        IO.puts("JSON output written to #{output_path}")
        IO.puts("Preview:")
        IO.puts(File.read!(output_path))

      {:error, reason} ->
        IO.puts(:stderr, "CLI failed: #{inspect(reason)}")
    end
  end

  defp maybe_add_base_url(args, nil), do: args
  defp maybe_add_base_url(args, ""), do: args
  defp maybe_add_base_url(args, url), do: args ++ ["--base-url", url]
end

Tinkex.Examples.CLIRunPromptFile.run()
</file>

<file path="cli_run_text.exs">
defmodule Tinkex.Examples.CLIRunText do
  @moduledoc false
  alias Tinkex.CLI

  def run do
    {:ok, _} = Application.ensure_all_started(:tinkex)

    api_key =
      System.get_env("TINKER_API_KEY") ||
        raise "Set TINKER_API_KEY to run this example"

    base_model = System.get_env("TINKER_BASE_MODEL", "meta-llama/Llama-3.1-8B")
    base_url = System.get_env("TINKER_BASE_URL")
    prompt = System.get_env("TINKER_PROMPT", "Hello from the CLI runner")

    max_tokens = env_integer("TINKER_MAX_TOKENS", 64)
    temperature = env_float("TINKER_TEMPERATURE", 0.7)
    num_samples = env_integer("TINKER_NUM_SAMPLES", 1)

    args =
      [
        "run",
        "--base-model",
        base_model,
        "--prompt",
        prompt,
        "--max-tokens",
        Integer.to_string(max_tokens),
        "--temperature",
        Float.to_string(temperature),
        "--num-samples",
        Integer.to_string(num_samples),
        "--api-key",
        api_key
      ]
      |> maybe_add_base_url(base_url)

    IO.puts("Running CLI with args: #{Enum.join(args, " ")}")

    case CLI.run(args) do
      {:ok, %{response: response}} ->
        IO.inspect(response, label: "sampling response")

      {:error, reason} ->
        IO.puts(:stderr, "CLI failed: #{inspect(reason)}")
    end
  end

  defp env_integer(var, default) do
    case System.get_env(var) do
      nil ->
        default

      value ->
        try do
          String.to_integer(value)
        rescue
          _ -> default
        end
    end
  end

  defp env_float(var, default) do
    case System.get_env(var) do
      nil ->
        default

      value ->
        try do
          String.to_float(value)
        rescue
          _ -> default
        end
    end
  end

  defp maybe_add_base_url(args, nil), do: args
  defp maybe_add_base_url(args, ""), do: args
  defp maybe_add_base_url(args, url), do: args ++ ["--base-url", url]
end

Tinkex.Examples.CLIRunText.run()
</file>

<file path="sampling_basic.exs">
defmodule Tinkex.Examples.SamplingBasic do
  @moduledoc false

  alias Tinkex.Error

  def run do
    {:ok, _} = Application.ensure_all_started(:tinkex)

    api_key =
      System.get_env("TINKER_API_KEY") ||
        raise "Set TINKER_API_KEY to run this example"

    base_url =
      System.get_env(
        "TINKER_BASE_URL",
        "https://tinker.thinkingmachines.dev/services/tinker-prod"
      )

    base_model = System.get_env("TINKER_BASE_MODEL", "meta-llama/Llama-3.1-8B")
    prompt_text = System.get_env("TINKER_PROMPT", "Hello from Tinkex!")

    max_tokens = env_integer("TINKER_MAX_TOKENS", 64)
    temperature = env_float("TINKER_TEMPERATURE", 0.7)
    num_samples = env_integer("TINKER_NUM_SAMPLES", 1)
    await_timeout = env_integer("TINKER_SAMPLE_TIMEOUT", 30_000)

    config = Tinkex.Config.new(api_key: api_key, base_url: base_url)
    {:ok, service} = Tinkex.ServiceClient.start_link(config: config)
    {:ok, sampler} = Tinkex.ServiceClient.create_sampling_client(service, base_model: base_model)

    {:ok, prompt} = Tinkex.Types.ModelInput.from_text(prompt_text, model_name: base_model)
    params = %Tinkex.Types.SamplingParams{max_tokens: max_tokens, temperature: temperature}

    IO.puts("Sampling #{num_samples} sequence(s) from #{base_model} ...")

    {:ok, task} =
      Tinkex.SamplingClient.sample(sampler, prompt, params,
        num_samples: num_samples,
        prompt_logprobs: false
      )

    case Task.await(task, await_timeout) do
      {:ok, response} ->
        IO.puts("Received #{length(response.sequences)} sequence(s):")

        Enum.with_index(response.sequences, 1)
        |> Enum.each(fn {seq, idx} ->
          text = decode(seq.tokens, base_model)
          IO.puts("Sample #{idx}: #{text}")
        end)

      {:error, error} ->
        IO.puts(:stderr, "Sampling failed: #{format_error(error)}")
    end
  end

  defp decode(tokens, model_name) do
    case Tinkex.Tokenizer.decode(tokens, model_name) do
      {:ok, text} -> text
      {:error, reason} -> "[decode failed: #{inspect(reason)}]"
    end
  end

  defp env_integer(var, default) do
    case System.get_env(var) do
      nil ->
        default

      value ->
        try do
          String.to_integer(value)
        rescue
          _ -> default
        end
    end
  end

  defp env_float(var, default) do
    case System.get_env(var) do
      nil ->
        default

      value ->
        try do
          String.to_float(value)
        rescue
          _ -> default
        end
    end
  end

  defp format_error(%Error{} = error), do: Error.format(error)
  defp format_error(other), do: inspect(other)
end

Tinkex.Examples.SamplingBasic.run()
</file>

<file path="sessions_management.exs">
defmodule Tinkex.Examples.SessionsManagement do
  @moduledoc """
  Example demonstrating session management APIs.

  Shows how to:
  - Create a RestClient
  - List all sessions
  - Get session details
  """

  alias Tinkex.{ServiceClient, RestClient, Config}

  def run do
    IO.puts("=== Tinkex Session Management Example ===\n")

    # Ensure application is started
    {:ok, _} = Application.ensure_all_started(:tinkex)

    # Create config from environment
    base_url =
      System.get_env("TINKER_BASE_URL") ||
        Application.get_env(
          :tinkex,
          :base_url,
          "https://tinker.thinkingmachines.dev/services/tinker-prod"
        )

    config =
      Config.new(
        api_key: System.get_env("TINKER_API_KEY") || raise("TINKER_API_KEY required"),
        base_url: base_url
      )

    # Start ServiceClient
    IO.puts("Starting ServiceClient...")
    {:ok, service_pid} = ServiceClient.start_link(config: config)

    # Create RestClient
    IO.puts("Creating RestClient...")
    {:ok, rest_client} = ServiceClient.create_rest_client(service_pid)

    # List sessions
    IO.puts("\n--- Listing Sessions ---")

    case RestClient.list_sessions(rest_client, limit: 10) do
      {:ok, response} ->
        IO.puts("Found #{length(response.sessions)} sessions:")

        Enum.each(response.sessions, fn session_id ->
          IO.puts("  • #{session_id}")
        end)

        # Get details for first session if available
        if length(response.sessions) > 0 do
          [first_session | _] = response.sessions
          get_session_details(rest_client, first_session)
        end

      {:error, error} ->
        IO.puts("Error listing sessions: #{inspect(error)}")
    end

    # Cleanup
    GenServer.stop(service_pid)
    IO.puts("\n=== Example Complete ===")
  end

  defp get_session_details(rest_client, session_id) do
    IO.puts("\n--- Session Details: #{session_id} ---")

    case RestClient.get_session(rest_client, session_id) do
      {:ok, response} ->
        IO.puts("Training Runs: #{length(response.training_run_ids)}")

        Enum.each(response.training_run_ids, fn run_id ->
          IO.puts("  • #{run_id}")
        end)

        IO.puts("Samplers: #{length(response.sampler_ids)}")

        Enum.each(response.sampler_ids, fn sampler_id ->
          IO.puts("  • #{sampler_id}")
        end)

      {:error, error} ->
        IO.puts("Error getting session: #{inspect(error)}")
    end
  end
end

Tinkex.Examples.SessionsManagement.run()
</file>

<file path="training_loop.exs">
defmodule Tinkex.Examples.TrainingLoop do
  @moduledoc false

  @default_base_url "https://tinker.thinkingmachines.dev/services/tinker-prod"
  @default_model "meta-llama/Llama-3.1-8B"
  @await_timeout 60_000

  alias Tinkex.Error
  alias Tinkex.Types.TensorData

  def run do
    {:ok, _} = Application.ensure_all_started(:tinkex)

    api_key = fetch_env!("TINKER_API_KEY")
    base_url = System.get_env("TINKER_BASE_URL", @default_base_url)
    base_model = System.get_env("TINKER_BASE_MODEL", @default_model)
    prompt = System.get_env("TINKER_PROMPT", "Fine-tuning sample prompt")
    sample_after? = System.get_env("TINKER_SAMPLE_AFTER_TRAIN", "0") not in ["0", "false", nil]
    sample_prompt = System.get_env("TINKER_SAMPLE_PROMPT", "Hello from fine-tuned weights!")

    IO.puts("Base URL: #{base_url}")
    IO.puts("Base model: #{base_model}")

    config = Tinkex.Config.new(api_key: api_key, base_url: base_url)

    with {:ok, service} <- Tinkex.ServiceClient.start_link(config: config),
         {:ok, training} <- create_training_client(service, base_model),
         {:ok, model_input} <- build_model_input(prompt, base_model, training) do
      run_training_steps(service, training, model_input, base_model, sample_after?, sample_prompt)
    else
      {:error, %Error{} = error} ->
        halt_with_error("Initialization failed", error)

      {:error, other} ->
        halt("Initialization failed: #{inspect(other)}")
    end
  end

  defp create_training_client(service, base_model) do
    Tinkex.ServiceClient.create_lora_training_client(service,
      base_model: base_model,
      lora_config: %Tinkex.Types.LoraConfig{rank: 16}
    )
  end

  defp build_model_input(prompt, base_model, training) do
    Tinkex.Types.ModelInput.from_text(prompt, model_name: base_model, training_client: training)
  end

  defp run_training_steps(
         service,
         training,
         model_input,
         base_model,
         sample_after?,
         sample_prompt
       ) do
    target_tokens = first_chunk_tokens(model_input)

    datum =
      Tinkex.Types.Datum.new(%{
        model_input: model_input,
        # Server expects loss_fn_inputs values as tensors with dtype/shape.
        loss_fn_inputs: %{
          target_tokens: to_tensor(target_tokens, :int64),
          weights: to_tensor(List.duplicate(1.0, length(target_tokens)), :float32)
        }
      })

    loop_start = System.monotonic_time(:millisecond)

    fb_task =
      start_task(
        Tinkex.TrainingClient.forward_backward(training, [datum], :cross_entropy),
        "forward_backward"
      )

    fb_output = await_task(fb_task, "forward_backward")
    IO.inspect(fb_output.metrics, label: "forward_backward metrics")

    optim_task =
      start_task(
        Tinkex.TrainingClient.optim_step(training, %Tinkex.Types.AdamParams{}),
        "optim_step"
      )

    optim_output = await_task(optim_task, "optim_step")
    IO.inspect(optim_output.metrics, label: "optim_step metrics")

    save_task =
      start_task(
        Tinkex.TrainingClient.save_weights_for_sampler(training),
        "save_weights_for_sampler"
      )

    save_result = await_task(save_task, "save_weights_for_sampler")
    IO.inspect(save_result, label: "save_weights_for_sampler response")

    if sample_after? do
      sample_with_saved_weights(service, save_result, base_model, sample_prompt)
    end

    duration_ms = System.monotonic_time(:millisecond) - loop_start
    IO.puts("Training loop finished in #{duration_ms} ms")
  end

  defp start_task(result, label) do
    case result do
      {:ok, task} ->
        task

      {:error, %Error{} = error} ->
        halt_with_error("#{label} failed", error)

      other ->
        halt("#{label} failed: #{inspect(other)}")
    end
  end

  defp await_task(task, label) do
    try do
      case Task.await(task, @await_timeout) do
        {:ok, result} ->
          result

        {:error, %Error{} = error} ->
          halt_with_error("#{label} error", error)

        other ->
          halt("#{label} returned unexpected response: #{inspect(other)}")
      end
    catch
      :exit, reason ->
        halt("#{label} task exited: #{inspect(reason)}")
    end
  end

  defp fetch_env!(var) do
    case System.get_env(var) do
      nil -> halt("Set #{var} to run this example")
      value -> value
    end
  end

  defp halt_with_error(prefix, %Error{} = error) do
    IO.puts(:stderr, "#{prefix}: #{Error.format(error)}")
    if error.data, do: IO.puts(:stderr, "Error data: #{inspect(error.data)}")
    System.halt(1)
  end

  defp halt(message) do
    IO.puts(:stderr, message)
    System.halt(1)
  end

  defp first_chunk_tokens(%Tinkex.Types.ModelInput{chunks: [chunk | _]}) do
    Map.get(chunk, :tokens) || Map.get(chunk, "tokens") || []
  end

  defp first_chunk_tokens(_), do: []

  defp to_tensor(tokens, dtype) when is_list(tokens) do
    seq_len = length(tokens)
    %TensorData{data: tokens, dtype: dtype, shape: [seq_len]}
  end

  defp to_tensor(_, dtype), do: %TensorData{data: [], dtype: dtype, shape: [0]}

  defp sample_with_saved_weights(service, save_result, base_model, sample_prompt)
       when is_map(save_result) do
    model_path = save_result["path"] || save_result[:path]
    sampling_session_id = save_result["sampling_session_id"] || save_result[:sampling_session_id]

    case create_sampler(service, model_path, sampling_session_id, base_model) do
      {:ok, sampler} ->
        do_sample(sampler, base_model, sample_prompt)

      {:error, reason} ->
        IO.puts(:stderr, "Skipping sampling; failed to create sampler: #{inspect(reason)}")
    end
  end

  defp sample_with_saved_weights(_service, _save_result, _base_model, _prompt), do: :ok

  defp create_sampler(service, model_path, _sampling_session_id, base_model) do
    opts =
      if model_path do
        [model_path: model_path, base_model: base_model]
      else
        # Fall back to base_model-only sampler if no path returned.
        [base_model: base_model]
      end
      |> maybe_put_sampling_session_seq_id()

    case Tinkex.ServiceClient.create_sampling_client(service, opts) do
      {:ok, sampler} -> {:ok, sampler}
      {:error, _} = error -> error
    end
  end

  defp maybe_put_sampling_session_seq_id(opts) do
    # Use first sampler seq_id for simplicity.
    Keyword.put_new(opts, :sampling_session_seq_id, 0)
  end

  defp do_sample(sampler, base_model, prompt) do
    {:ok, model_input} = Tinkex.Types.ModelInput.from_text(prompt, model_name: base_model)
    params = %Tinkex.Types.SamplingParams{max_tokens: 32, temperature: 0.7}

    with {:ok, task} <-
           Tinkex.SamplingClient.sample(sampler, model_input, params,
             num_samples: 1,
             prompt_logprobs: false
           ),
         {:ok, response} <- Task.await(task, 30_000) do
      [seq | _] = response.sequences
      text = decode(seq.tokens, base_model)
      IO.puts("Sample from saved weights: #{text}")
    else
      {:error, error} ->
        IO.puts(:stderr, "Sampling after train failed: #{inspect(error)}")
    end
  end

  defp decode(tokens, model_name) do
    case Tinkex.Tokenizer.decode(tokens, model_name) do
      {:ok, text} -> text
      {:error, reason} -> "[decode failed: #{inspect(reason)}]"
    end
  end
end

Tinkex.Examples.TrainingLoop.run()
</file>

</files>