ruby-laya 0.1.0

data/lib/laya/agent.rb

@@ -0,0 +1,324 @@
+# frozen_string_literal: true
+
+require "json"
+require_relative "runtime"
+require_relative "tokenizer"
+require_relative "hub"
+require_relative "question"
+require_relative "result"
+
+module Laya
+  # A Laya checkpoint, loaded and ready to answer typed questions.
+  #
+  #   agent = Laya.load("convaiinnovations/laya")
+  #   result = agent.predict({ "message" => "I was charged twice" }, Laya.triage_questions)
+  #   result[:intent].choice        # => "refund"
+  #   result[:churn_risk].probability
+  #
+  # Every question in a call is answered in one forward pass. The weights are the ones Convai
+  # Innovations published, exported to ONNX; see {Checkpoints}.
+  class Agent
+    # What a checkpoint directory must hold; see {Laya::Checkpoints::RUNTIME_FILES}.
+    RUNTIME_FILES = Checkpoints::RUNTIME_FILES
+
+    # The traced graph fixes the branch upstream picks at runtime for a one-option question, so a
+    # batch always carries at least two markers; the spare is masked off and scores nothing.
+    MIN_MARKERS = 2
+
+    # The shortest sequence the graph accepts. A batch of very short texts is padded up to it,
+    # which changes nothing: the padding is masked off.
+    MIN_SEQ = 8
+
+    DEFAULT_MAX_LEN = 512
+    DEFAULT_HEAD_MAX_LEN = 192
+
+    attr_reader :model_id, :model_dir, :config, :onnx_config, :tokenizer, :runtime,
+                :temperature, :temperature_by_options, :temperature_raw, :temperature_by_options_raw
+
+    class << self
+      # Load a checkpoint. See {Laya.load}, which is the documented entry point.
+      def load(model_id_or_path = Checkpoints::BUNDLE_REPO, **, &block)
+        agent = new(model_id_or_path, **)
+        return agent unless block
+
+        begin
+          block.call(agent)
+        ensure
+          agent.close
+        end
+      end
+
+      # The local directory holding `model_id_or_path`, downloading it when it names a repository.
+      #
+      # `hub` is the downloader, injectable so an app can serve exports from its own store: it
+      # answers `snapshot(repo, subfolder:, allow_patterns:, token:, revision:)` with a local path.
+      def resolve_dir(model_id_or_path, subfolder: nil, token: nil, revision: Hub::DEFAULT_REVISION,
+                      hub: Hub)
+        return local_dir(model_id_or_path, subfolder) if File.directory?(model_id_or_path.to_s)
+
+        if path_like?(model_id_or_path)
+          raise ModelNotFoundError,
+                "local model path not found: #{model_id_or_path.inspect}. Check the directory exists " \
+                "and that the export wrote it successfully."
+        end
+
+        repo, onnx_subfolder = onnx_source_for(model_id_or_path, subfolder)
+        hub.snapshot(repo, subfolder: onnx_subfolder, allow_patterns: RUNTIME_FILES,
+                           token: token, revision: revision)
+      end
+
+      # The ONNX repository and subfolder serving an upstream model id, or the id itself when it
+      # already names an export.
+      def onnx_source_for(model_id, subfolder)
+        name = Checkpoints.resolve_upstream(model_id, subfolder)
+        return Checkpoints.source_for(name) if name
+
+        [model_id.to_s, subfolder]
+      end
+
+      def path_like?(model_id_or_path)
+        value = model_id_or_path.to_s
+        value.start_with?("/", "./", "../", "~") || File.absolute_path?(value)
+      end
+
+      def local_dir(path, subfolder)
+        dir = subfolder ? File.join(path, subfolder) : path.to_s
+        return dir if File.directory?(dir)
+
+        raise ModelNotFoundError, "subfolder #{subfolder.inspect} not found in #{path.inspect}"
+      end
+    end
+
+    # @param model_id_or_path [String] an upstream model id, an ONNX repository id, or a directory
+    # @param device [String, Symbol, nil] "cpu" (default), "coreml", "cuda", ...
+    # @param providers [Array<String>, nil] ONNX Runtime providers, overriding `device`
+    # @param threads [Integer, nil] intra-op threads; ONNX Runtime decides when nil
+    # @param hub [#snapshot] where repository ids are downloaded from; defaults to {Laya::Hub}
+    def initialize(model_id_or_path = Checkpoints::BUNDLE_REPO, device: nil, providers: nil,
+                   token: nil, subfolder: nil, threads: nil, revision: Hub::DEFAULT_REVISION,
+                   hub: Hub)
+      @model_id = model_id_or_path.to_s
+      @model_dir = Agent.resolve_dir(model_id_or_path, subfolder: subfolder, token: token,
+                                                       revision: revision, hub: hub)
+      @config = read_json("rl_agent_config.json")
+      @onnx_config = read_json("onnx_config.json")
+      @tokenizer = Tokenizer.from_dir(File.join(@model_dir, "tokenizer"))
+      @temperature_raw = config.fetch("temperature", [1.0, 1.0, 1.0])
+      @temperature_by_options_raw = config.fetch("temperature_by_options", {})
+      @temperature = @temperature_raw.map { |value| Common.clamp_temperature(value) }
+      @temperature_by_options = @temperature_by_options_raw.transform_values { |v| Common.clamp_temperature(v) }
+      warn_about_temperatures
+      @runtime = Runtime.new(onnx_path, providers: providers, device: device, threads: threads)
+    end
+
+    # Evaluate every question against `state` in a single forward pass.
+    #
+    # `state` is a String, a Hash or an Array of conversation turns. `questions` maps a question id
+    # to its definition; see {Question}. Returns a {Result}.
+    def predict(state, questions)
+      unless questions.is_a?(Hash)
+        raise ArgumentError, "questions must be a Hash of question id => definition, got #{questions.class}"
+      end
+      return Result.new(answers: {}, usage: EMPTY_USAGE.dup) if questions.empty?
+
+      asked = questions.map { |id, definition| Question.build(id, definition) }
+      batch = tokenize(state, asked)
+      logits, act_logits = runtime.decide(batch)
+      Result.new(answers: collect_answers(asked, batch, logits, act_logits),
+                 usage: { "input_tokens" => batch[:input_tokens], "output_tokens" => 0 })
+    end
+    alias system_one predict
+
+    # Mean-pooled encoder states for `texts`, one row per string.
+    #
+    # This is what {Laya.embed_fn_from_agent} hands the shortlist. Padding is excluded from the
+    # mean, and the decision head never runs.
+    def embed(texts, max_length: nil, batch_size: 32)
+      rows = Array(texts).map { |text| text.nil? ? "" : text.to_s }
+      return [] if rows.empty?
+
+      max_length ||= max_len
+      rows.each_slice(batch_size).flat_map do |chunk|
+        encoded = pad_to_min_seq(tokenizer.encode_batch(chunk, max_length: max_length))
+        hidden = runtime.encode(embedding_batch(encoded))
+        mean_pool(hidden, encoded["attention_mask"])
+      end
+    end
+
+    # The token budget for the whole sequence, and the share of it the options may use.
+    def max_len
+      config.fetch("max_len", DEFAULT_MAX_LEN)
+    end
+
+    def head_max_len
+      config.fetch("head_max_len", DEFAULT_HEAD_MAX_LEN)
+    end
+
+    def hidden_size
+      onnx_config["hidden_size"]
+    end
+
+    # Release the ONNX session. The agent cannot answer afterwards.
+    def close
+      runtime.close
+      self
+    end
+
+    def closed?
+      runtime.closed?
+    end
+
+    def inspect
+      "#<Laya::Agent #{model_id.inspect} providers=#{runtime.providers.inspect}>"
+    end
+
+    private
+
+    EMPTY_USAGE = { "input_tokens" => 0, "output_tokens" => 0 }.freeze
+    private_constant :EMPTY_USAGE
+
+    def onnx_path
+      path = File.join(model_dir, "model.onnx")
+      return path if File.file?(path)
+
+      raise IncompatibleModelError,
+            "no 'model.onnx' in #{model_dir}. ruby-laya runs the ONNX exports of the Laya " \
+            "checkpoints; export one with tools/export_onnx.py or load #{Checkpoints.onnx_repo}."
+    end
+
+    def read_json(name)
+      path = File.join(model_dir, name)
+      unless File.file?(path)
+        raise IncompatibleModelError,
+              "no #{name.inspect} in #{model_dir}. That file ships with an exported Laya " \
+              "checkpoint, so load one of those or re-export with tools/export_onnx.py."
+      end
+
+      JSON.parse(File.read(path))
+    rescue JSON::ParserError => e
+      raise IncompatibleModelError, "#{name} in #{model_dir} is not valid JSON: #{e.message}"
+    end
+
+    # Tokenize every question into one padded batch of marker sequences.
+    def tokenize(state, asked)
+      items = asked.map do |question|
+        ids, markers = Common.build_sequence(tokenizer, state, question.internal,
+                                             max_len: max_len, head_max_len: head_max_len)
+        if markers.length != question.options.length
+          raise ArgumentError,
+                "question #{question.id.inspect} has #{question.options.length} options, which do not " \
+                "fit head_max_len=#{head_max_len}; shortlist them or raise the budget"
+        end
+
+        { ids: ids, markers: markers, qtype: question.qtype }
+      end
+      collate(items)
+    end
+
+    def collate(items)
+      width = [items.map { |item| item[:ids].length }.max, min_seq].max
+      markers = [items.map { |item| item[:markers].length }.max, MIN_MARKERS].max
+      pad = tokenizer.pad_token_id
+      {
+        input_ids: items.map { |item| item[:ids] + Array.new(width - item[:ids].length, pad) },
+        attention_mask: items.map { |item| Array.new(item[:ids].length, 1) + Array.new(width - item[:ids].length, 0) },
+        marker_pos: items.map { |item| item[:markers] + Array.new(markers - item[:markers].length, 0) },
+        marker_mask: items.map do |item|
+          Array.new(item[:markers].length, true) + Array.new(markers - item[:markers].length, false)
+        end,
+        qtype: items.map { |item| item[:qtype] },
+        input_tokens: items.sum { |item| item[:ids].length },
+        markers: items.map { |item| item[:markers].length }
+      }
+    end
+
+    def collect_answers(asked, batch, logits, act_logits)
+      act = act_logits.map { |row| Common.softmax(row) }
+      asked.each_with_index.to_h do |question, row|
+        options = batch[:markers][row]
+        scale = temperature_for(question, options)
+        probabilities = Common.softmax(logits[row].first(options), temperature: scale)
+        [question.id, question.answer(
+          probabilities: probabilities,
+          confidence: Common.confidence_from_probs(probabilities, options).round(4),
+          action_probability: act[row][0].round(4)
+        )]
+      end
+    end
+
+    # The fitted temperature for this question type and option count, clamped at load.
+    def temperature_for(question, options)
+      temperature_by_options.fetch(Common.temp_bucket(question.qtype, options)) do
+        temperature[question.qtype]
+      end
+    end
+
+    # The graph's shortest accepted sequence, as the export recorded it.
+    def min_seq
+      onnx_config.fetch("min_seq", MIN_SEQ)
+    end
+
+    # Short texts, and texts that tokenize to nothing at all, still have to form a valid batch.
+    def pad_to_min_seq(encoded)
+      width = encoded["input_ids"].map(&:length).max.to_i
+      return encoded if width >= min_seq
+
+      pad = tokenizer.pad_token_id
+      { "input_ids" => encoded["input_ids"].map { |row| row + Array.new(min_seq - row.length, pad) },
+        "attention_mask" => encoded["attention_mask"].map { |row| row + Array.new(min_seq - row.length, 0) } }
+    end
+
+    def embedding_batch(encoded)
+      rows = encoded["input_ids"].length
+      {
+        input_ids: encoded["input_ids"],
+        attention_mask: encoded["attention_mask"],
+        marker_pos: Array.new(rows) { Array.new(MIN_MARKERS, 0) },
+        marker_mask: Array.new(rows) { Array.new(MIN_MARKERS, true) },
+        qtype: Array.new(rows, 0)
+      }
+    end
+
+    def mean_pool(hidden, attention_mask)
+      hidden.each_with_index.map do |sequence, row|
+        mask = attention_mask[row]
+        live = mask.sum
+        next Array.new(sequence.first.length, 0.0) if live.zero?
+
+        sums = Array.new(sequence.first.length, 0.0)
+        sequence.each_with_index do |vector, position|
+          next if mask[position].zero?
+
+          vector.each_with_index { |value, i| sums[i] += value }
+        end
+        sums.map { |value| value / live }
+      end
+    end
+
+    # A checkpoint may ship a temperature that sharpens rather than softens; those are clamped at
+    # load, and saying so once is the only warning a caller gets.
+    def warn_about_temperatures
+      entries = @temperature_by_options_raw.map { |bucket, raw| [bucket, raw, @temperature_by_options[bucket]] }
+      entries += @temperature_raw.each_with_index.map { |raw, i| ["temperature[#{i}]", raw, @temperature[i]] }
+      rejected = entries.filter_map do |name, raw, applied|
+        format("%s=%s -> %g", name, raw.inspect, applied) if clamped?(raw, applied)
+      end
+      return if rejected.empty?
+
+      warn "[laya] #{model_id}: this checkpoint ships invalid temperatures or values outside " \
+           "[#{TEMP_MIN}, #{TEMP_MAX}]; using #{rejected.join(', ')}. Treat confidence from the " \
+           "affected entries as uncalibrated."
+    end
+
+    # True when the value actually applied is not the one the checkpoint shipped. An exact
+    # comparison is what is wanted here: the clamp either returned the same number or a
+    # different one.
+    def clamped?(raw, applied)
+      !Float(raw).equal?(applied) && Float(raw) != applied # rubocop:disable Lint/FloatComparison
+    rescue ArgumentError, TypeError
+      true
+    end
+  end
+
+  RLAgent = Agent
+end

data/lib/laya/ask.rb

@@ -0,0 +1,73 @@
+# frozen_string_literal: true
+
+require_relative "questions"
+
+module Laya
+  # A decision built inline, for the questions not worth a class.
+  #
+  #   answers = Laya.ask(email)
+  #                 .choice(:department, "Which team?", billing: "invoices", technical: "outages")
+  #                 .noul(:refund, "Do they want money back?")
+  #                 .decide
+  #
+  #   answers.department == :billing # => true
+  #   answers[:refund].probability   # => 0.856
+  #
+  # Each call returns the builder, so the chain reads as the question set it is. {#decide} runs
+  # them all in one forward pass and hands back the {Result}.
+  class Ask
+    attr_reader :state, :questions
+
+    def initialize(state, client: nil, **options)
+      @state = state
+      @client = client
+      @options = options
+      @questions = {}
+    end
+
+    def choice(name, instructions, criteria = nil, **labels)
+      add(name, Questions.choice(instructions, criteria, **labels))
+    end
+
+    def score(name, instructions, levels:)
+      add(name, Questions.score(instructions, levels))
+    end
+
+    def noul(name, instructions, yes: nil, no: nil)
+      add(name, Questions.noul(instructions, yes: yes, no: no))
+    end
+
+    # Send it to a specific checkpoint rather than letting the router choose.
+    def using(model)
+      @options = @options.merge(model: model)
+      self
+    end
+
+    # Answer every question asked so far, in one forward pass.
+    def decide
+      raise ArgumentError, "ask at least one question before calling decide" if questions.empty?
+
+      client.predict(state, questions, **runnable_options)
+    end
+
+    def inspect
+      "#<Laya::Ask #{questions.keys.inspect}>"
+    end
+
+    private
+
+    def add(name, question)
+      @questions[name] = question
+      self
+    end
+
+    def client
+      @client || Laya.client
+    end
+
+    # `model:` only means something to a Router; an Agent already is one checkpoint.
+    def runnable_options
+      client.is_a?(Router) ? @options : @options.except(:model)
+    end
+  end
+end

data/lib/laya/checkpoints.rb

@@ -0,0 +1,77 @@
+# frozen_string_literal: true
+
+module Laya
+  # Where the three Laya checkpoints come from.
+  #
+  # Upstream publishes PyTorch weights; this gem runs ONNX exports of those same weights, so every
+  # upstream model id is mapped to its export here. `Laya.load("convaiinnovations/laya")` therefore
+  # works as the Python docs describe, and resolves to the English export.
+  module Checkpoints
+    # The repository holding the exports. Override with `LAYA_ONNX_REPO` to serve your own, for
+    # example a private mirror or a re-export of a fine-tuned checkpoint.
+    def self.onnx_repo
+      ENV.fetch("LAYA_ONNX_REPO", "codenamev/laya-onnx")
+    end
+
+    # The bundle repo upstream publishes, and the standalone repo for each checkpoint.
+    BUNDLE_REPO = "convaiinnovations/laya"
+
+    NAMES = %w[english multilingual typed-decisions].freeze
+
+    # What an exported checkpoint consists of, and so what the gem downloads. Here rather than on
+    # the runtime, because it describes the files, not how they are executed.
+    RUNTIME_FILES = ["model.onnx", "onnx_config.json", "rl_agent_config.json", "tokenizer/*"].freeze
+
+    # name => [upstream bundle subfolder, standalone upstream repo, ONNX subfolder]
+    SOURCES = {
+      "english" => [nil, "convaiinnovations/laya", "english"],
+      "multilingual" => ["multilingual", "convaiinnovations/laya-multilingual", "multilingual"],
+      "typed-decisions" => ["typed-decisions", "convaiinnovations/laya-typed-decisions", "typed-decisions"]
+    }.freeze
+
+    # Aliases people are likely to type.
+    ALIASES = {
+      "en" => "english", "laya" => "english", "default" => "english",
+      "multi" => "multilingual", "ml" => "multilingual", "laya-multilingual" => "multilingual",
+      "typed" => "typed-decisions", "typed_decisions" => "typed-decisions",
+      "laya-typed-decisions" => "typed-decisions", "decisions" => "typed-decisions"
+    }.freeze
+
+    module_function
+
+    # The canonical checkpoint name for `name` or one of its aliases.
+    def normalise(name)
+      key = name.to_s.strip.downcase
+      key = ALIASES.fetch(key, key)
+      return key if SOURCES.key?(key)
+
+      raise ArgumentError, "unknown model #{name.inspect}; choose one of #{NAMES} " \
+                           "(or an alias: #{ALIASES.keys.sort})"
+    end
+
+    # The ONNX source for a checkpoint name, as `[repo, subfolder]`.
+    def source_for(name)
+      [onnx_repo, SOURCES.fetch(normalise(name))[2]]
+    end
+
+    # The checkpoint an upstream model id names, or nil when it names none.
+    #
+    #   resolve_upstream("convaiinnovations/laya")                     # => "english"
+    #   resolve_upstream("convaiinnovations/laya", "multilingual")     # => "multilingual"
+    #   resolve_upstream("convaiinnovations/laya-multilingual")        # => "multilingual"
+    def resolve_upstream(repo, subfolder = nil)
+      repo = repo.to_s
+      SOURCES.each do |name, (bundle_subfolder, standalone, _onnx)|
+        return name if repo == BUNDLE_REPO && subfolder.to_s == bundle_subfolder.to_s
+        return name if repo == standalone && subfolder.nil?
+      end
+      nil
+    end
+
+    # Human-readable id for a source spec: "repo" or "repo/subfolder".
+    def repo_str(spec)
+      repo, subfolder = Array(spec)
+      subfolder ? "#{repo}/#{subfolder}" : repo.to_s
+    end
+  end
+end

data/lib/laya/common.rb

@@ -0,0 +1,175 @@
+# frozen_string_literal: true
+
+module Laya
+  # Question type => the index the decision head's type embedding uses.
+  QTYPES = { "choice" => 0, "score" => 1, "noul" => 2 }.freeze
+  QTYPE_NAMES = QTYPES.invert.freeze
+
+  # A fitted temperature below 1 sharpens the logits instead of softening them. One shipped
+  # bucket is 0.1006, which multiplies them tenfold: a 0.24 top probability is published as 0.99,
+  # so a caller gating on confidence is told a coin flip is a certainty. No honest calibration
+  # needs to sharpen this hard, so values outside these bounds are refused.
+  TEMP_MIN = 0.5
+  TEMP_MAX = 5.0
+
+  # Token sequence construction, option rendering and the calibration arithmetic, shared by the
+  # runtime, the router and the shortlist. Pure Ruby, and a faithful port: the checkpoints were
+  # trained on exactly these strings.
+  module Common
+    # An option's text may not run past this many tokens.
+    MAX_OPTION_TOKENS = 48
+
+    module_function
+
+    # The text form of a state: a String passes through, anything else becomes Python-style JSON.
+    def serialize_state(state)
+      return state if state.is_a?(String)
+
+      PyJSON.dumps(state, default: :to_s)
+    end
+
+    # One criterion rendered as text. Strings pass through; anything structured becomes compact
+    # JSON, so a rubric reads as JSON rather than as a language's idea of `inspect`.
+    def render_criterion(value)
+      return value if value.is_a?(String)
+
+      PyJSON.dumps(value, default: :to_s)
+    end
+
+    # The option texts in label order. A noul question is always `[false, true]`.
+    def render_options(question)
+      type = Util.get(question, "t").to_s
+      criteria = Util.get(question, "crit")
+      case type
+      when "choice" then choice_options(criteria)
+      when "score" then criteria.each_with_index.map { |level, i| "level #{i}: #{render_criterion(level)}" }
+      else noul_options(criteria || {})
+      end
+    end
+
+    # Only nil and "" mean "no description"; 0 and false are legitimate criterion values.
+    def choice_options(criteria)
+      criteria.map do |label, description|
+        blank?(description) ? label.to_s : "#{label}: #{render_criterion(description)}"
+      end
+    end
+
+    def noul_options(criteria)
+      false_text = Util.get(criteria, "false")
+      true_text = Util.get(criteria, "true")
+      ["false: #{blank?(false_text) ? 'no, the statement does not hold' : render_criterion(false_text)}",
+       "true: #{blank?(true_text) ? 'yes, the statement holds' : render_criterion(true_text)}"]
+    end
+
+    def blank?(value)
+      value.nil? || value == ""
+    end
+
+    # Build one question's token sequence:
+    #
+    #   [CLS] <type> instructions [SEP] [MASK] opt0 [MASK] opt1 ... [SEP] state [SEP]
+    #
+    # Returns `[token_ids, marker_positions]`, where each marker is the `[MASK]` in front of an
+    # option and is what the decision head scores. `tokenizer` answers `mask_token`,
+    # `mask_token_id`, `cls_token_id`, `sep_token_id` and `encode_ids`.
+    def build_sequence(tokenizer, state, question, max_len: 512, head_max_len: 192,
+                       option_order: nil, truncate_left: false)
+      mask = tokenizer.mask_token
+      options = render_options(question)
+      order = option_order || (0...options.length).to_a
+      instructions = Util.get(question, "ins").to_s.gsub(mask, " ")
+
+      head = tokenizer.encode_ids("#{Util.get(question, 't')} question: #{instructions}")
+      rendered = order.map do |i|
+        [tokenizer.mask_token_id] + tokenizer.encode_ids(" #{options[i].gsub(mask, ' ')}").first(MAX_OPTION_TOKENS)
+      end
+      rendered = trim_options(rendered, head_max_len) if head_max_len - rendered.sum(&:length) < 16
+      head = head.first([8, head_max_len - rendered.sum(&:length)].max)
+
+      ids = [tokenizer.cls_token_id] + head + [tokenizer.sep_token_id]
+      markers = []
+      rendered.each do |option|
+        markers << ids.length
+        ids.concat(option)
+      end
+      ids << tokenizer.sep_token_id
+
+      room = [0, max_len - ids.length - 1].max
+      state_ids = tokenizer.encode_ids(serialize_state(state).gsub(mask, " "))
+      state_ids = truncate_left ? state_ids.last(room) : state_ids.first(room)
+      ids = ids + state_ids + [tokenizer.sep_token_id]
+      [ids.first(max_len), markers.select { |marker| marker < max_len }]
+    end
+
+    # Too many options for the budget: every one of them is cut to an equal share.
+    def trim_options(rendered, head_max_len)
+      per = [4, (head_max_len - 16) / [1, rendered.length].max].max
+      rendered.map { |option| option.first(per) }
+    end
+
+    # Confidence as normalized Shannon entropy: `1 - H(p) / log(k)`.
+    def confidence_from_probs(probabilities, k)
+      return 1.0 if k < 2
+
+      entropy = -probabilities.first(k).sum { |p| p * Math.log(p.clamp(1e-12, 1.0)) }
+      (1.0 - (entropy / Math.log(k))).clamp(0.0, 1.0)
+    end
+
+    # Expected Calibration Error across confidence bins. The first bin includes its lower edge,
+    # so a prediction of exactly zero confidence is counted rather than dropped.
+    def ece_score(confidences, correct, bins: 15)
+      return Float::NAN if confidences.empty?
+
+      correct = correct.map do |value|
+        if value == true
+          1.0
+        else
+          value == false ? 0.0 : value.to_f
+        end
+      end
+      total = confidences.length.to_f
+      (0...bins).sum do |bin|
+        low = bin.fdiv(bins)
+        high = (bin + 1).fdiv(bins)
+        selected = confidences.each_index.select do |i|
+          (bin.zero? ? confidences[i] >= low : confidences[i] > low) && confidences[i] <= high
+        end
+        next 0.0 if selected.empty?
+
+        mean_confidence = selected.sum { |i| confidences[i] } / selected.length
+        mean_correct = selected.sum { |i| correct[i] } / selected.length
+        (selected.length / total) * (mean_confidence - mean_correct).abs
+      end
+    end
+
+    # The calibration bucket for a question type (index or name) and option count.
+    def temp_bucket(qtype, k)
+      name = qtype.is_a?(Integer) ? QTYPE_NAMES.fetch(qtype) : qtype.to_s
+      size = if k <= 2 then "2"
+             elsif k <= 5 then "3-5"
+             elsif k <= 10 then "6-10"
+             else "11+"
+             end
+      "#{name}:#{size}"
+    end
+
+    # A usable temperature: `value` confined to the sane range, or 1.0 when it is not a number.
+    def clamp_temperature(value, low: TEMP_MIN, high: TEMP_MAX)
+      number = Float(value)
+      return 1.0 if number.nan? || number.infinite?
+
+      number.clamp(low, high)
+    rescue ArgumentError, TypeError
+      1.0
+    end
+
+    # Softmax over logits, optionally tempered.
+    def softmax(logits, temperature: 1.0)
+      scaled = logits.map { |logit| logit / temperature }
+      highest = scaled.max
+      exponentials = scaled.map { |value| Math.exp(value - highest) }
+      total = exponentials.sum
+      exponentials.map { |value| value / total }
+    end
+  end
+end