metrics.hpp

Go to the documentation of this file.

#pragma once
 
// SPDX-License-Identifier: Apache-2.0
// Copyright 2026 Lloyal Labs
 
#include <algorithm>
#include <cmath>
#include <cstdint>
#include <limits>
 
namespace lloyal::metrics {
 
// ============================================================================
// Types
// ============================================================================
 
enum class Base { Nats, Bits };
 
// ============================================================================
// Internal helpers (ported from metrics.ts)
// ============================================================================
 
namespace detail {
 
constexpr float LN2 = 0.693147180559945309417232121458176568f;
 
inline float max_finite(const float* a, int n) {
  float m = -std::numeric_limits<float>::infinity();
  for (int i = 0; i < n; ++i) {
    const float v = a[i];
    if (std::isfinite(v) && v > m) m = v;
  }
  return m;
}
 
inline float log_sum_exp(const float* a, int n, float shift) {
  float s = 0.0f;
  for (int i = 0; i < n; ++i) {
    const float v = a[i];
    if (std::isfinite(v)) s += std::exp(v - shift);
  }
  if (s == 0.0f) return -std::numeric_limits<float>::infinity();
  return shift + std::log(s);
}
 
}  // namespace detail
 
// ============================================================================
// Perplexity tracking types (used by BranchStore registry)
// ============================================================================
 
struct PerplexityState {
  float nll_sum_nats = 0.0f;
  int count = 0;
};
 
struct BranchMetricsState {
  PerplexityState model;    
  PerplexityState sampling; 
};
 
// ============================================================================
// Model-level metrics (raw logits, before filters)
// ============================================================================
 
inline float model_surprisal(
    const float* logits,
    int n_vocab,
    int picked_id,
    Base base = Base::Nats
) {
  if (!logits || n_vocab == 0) {
    return std::numeric_limits<float>::infinity();
  }
  if (picked_id < 0 || picked_id >= n_vocab) {
    return std::numeric_limits<float>::infinity();
  }
 
  const float picked = logits[picked_id];
  if (!std::isfinite(picked)) return std::numeric_limits<float>::infinity();
 
  const float m = detail::max_finite(logits, n_vocab);
  if (!std::isfinite(m)) return std::numeric_limits<float>::infinity();
 
  const float log_z = detail::log_sum_exp(logits, n_vocab, m);
  if (!std::isfinite(log_z)) return std::numeric_limits<float>::infinity();
 
  const float surprisal_nats = std::max(0.0f, -(picked - log_z));
  return base == Base::Bits ? surprisal_nats / detail::LN2 : surprisal_nats;
}
 
inline float model_entropy(
    const float* logits,
    int n_vocab,
    Base base = Base::Nats
) {
  if (!logits || n_vocab == 0) {
    return std::numeric_limits<float>::infinity();
  }
 
  const float m = detail::max_finite(logits, n_vocab);
  if (!std::isfinite(m)) return std::numeric_limits<float>::infinity();
 
  const float log_z = detail::log_sum_exp(logits, n_vocab, m);
  if (!std::isfinite(log_z)) return std::numeric_limits<float>::infinity();
 
  float ez = 0.0f;
  for (int i = 0; i < n_vocab; ++i) {
    const float z = logits[i];
    if (!std::isfinite(z)) continue;
    const float p = std::exp(z - log_z);
    ez += p * z;
  }
 
  const float h_nats = std::max(0.0f, log_z - ez);
  return base == Base::Bits ? h_nats / detail::LN2 : h_nats;
}
 
// ============================================================================
// Sampling-level metrics (post-filter logits, after top-k/p/temp)
// ============================================================================
 
inline float sampling_surprisal(
    const float* candidate_logits,
    const int32_t* candidate_ids,
    int n_candidates,
    int picked_id,
    Base base = Base::Nats
) {
  if (!candidate_logits || !candidate_ids || n_candidates == 0) {
    return std::numeric_limits<float>::infinity();
  }
 
  // Find picked_id in candidates
  int local = -1;
  for (int i = 0; i < n_candidates; ++i) {
    if (candidate_ids[i] == picked_id) {
      local = i;
      break;
    }
  }
  if (local == -1) return std::numeric_limits<float>::infinity();
  if (n_candidates == 1) return 0.0f;
 
  const float picked = candidate_logits[local];
  if (!std::isfinite(picked)) return std::numeric_limits<float>::infinity();
 
  const float m = detail::max_finite(candidate_logits, n_candidates);
  if (!std::isfinite(m)) return std::numeric_limits<float>::infinity();
 
  const float log_z = detail::log_sum_exp(candidate_logits, n_candidates, m);
  if (!std::isfinite(log_z)) return std::numeric_limits<float>::infinity();
 
  const float surprisal_nats = std::max(0.0f, -(picked - log_z));
  return base == Base::Bits ? surprisal_nats / detail::LN2 : surprisal_nats;
}
 
inline float sampling_entropy(
    const float* candidate_logits,
    int n_candidates,
    Base base = Base::Nats
) {
  if (!candidate_logits || n_candidates == 0) {
    return std::numeric_limits<float>::infinity();
  }
  if (n_candidates == 1) return 0.0f;
 
  const float m = detail::max_finite(candidate_logits, n_candidates);
  if (!std::isfinite(m)) return std::numeric_limits<float>::infinity();
 
  const float log_z = detail::log_sum_exp(candidate_logits, n_candidates, m);
  if (!std::isfinite(log_z)) return std::numeric_limits<float>::infinity();
 
  float ez = 0.0f;
  for (int i = 0; i < n_candidates; ++i) {
    const float z = candidate_logits[i];
    if (!std::isfinite(z)) continue;
    const float p = std::exp(z - log_z);
    ez += p * z;
  }
 
  const float h_nats = std::max(0.0f, log_z - ez);
  return base == Base::Bits ? h_nats / detail::LN2 : h_nats;
}
 
}  // namespace lloyal::metrics