Moonbit docs

///| Reproducible synthetic decoding workloads. They are not a replacement for

///| real model evaluation; their purpose is to exercise acceptance, rejection,

///|
/// vocabulary shape, and scheduling paths in CI without downloading weights.
pub enum WorkloadError {
  InvalidRounds(Int)
  InvalidDepth(Int)
  InvalidVocabulary(Int)
  InvalidAgreement(Double)
  RandomFailure
} derive(Eq, Debug)

///| Agreement controls how often target and draft favour the same token.

///| Values near one create easy speculative cases; values near zero stress

///|
/// residual sampling and rejection handling.
pub struct WorkloadConfig {
  rounds : Int
  proposal_depth : Int
  vocabulary_size : Int
  agreement : Double
}

///|
pub fn WorkloadConfig::new(
  rounds : Int,
  proposal_depth : Int,
  vocabulary_size : Int,
  agreement : Double,
) -> Result[WorkloadConfig, WorkloadError] {
  if rounds <= 0 {
    return Err(InvalidRounds(rounds))
  }
  if proposal_depth <= 0 {
    return Err(InvalidDepth(proposal_depth))
  }
  if vocabulary_size < 2 {
    return Err(InvalidVocabulary(vocabulary_size))
  }
  if !finite_number(agreement) || agreement < 0.0 || agreement > 1.0 {
    return Err(InvalidAgreement(agreement))
  }
  Ok({ rounds, proposal_depth, vocabulary_size, agreement })
}

///|
/// A small default workload appropriate for the command-line example.
pub fn WorkloadConfig::demo() -> WorkloadConfig {
  { rounds: 4, proposal_depth: 3, vocabulary_size: 8, agreement: 0.75 }
}

///|
pub fn WorkloadConfig::rounds(self : WorkloadConfig) -> Int {
  self.rounds
}

///|
pub fn WorkloadConfig::proposal_depth(self : WorkloadConfig) -> Int {
  self.proposal_depth
}

///|
pub fn WorkloadConfig::vocabulary_size(self : WorkloadConfig) -> Int {
  self.vocabulary_size
}

///| Build one noisy logit row with a chosen winning token. Adding a positive

///| margin keeps the scenario easy to reason about while background noise

///|
/// prevents every distribution from being a trivial one-hot vector.
fn logits_for(
  rng : DeterministicRng,
  vocabulary_size : Int,
  preferred : Int,
) -> Array[Double] {
  let logits : Array[Double] = []
  for token in 0.. Result[Array[SimulationRound], WorkloadError] {
  let rng = match DeterministicRng::new(seed) {
    Ok(value) => value
    Err(_) => return Err(RandomFailure)
  }
  let schedule : Array[SimulationRound] = []
  for _ in 0.. value
        Err(_) => return Err(RandomFailure)
      }
      let target_token = if rng.next_unit() < config.agreement {
        draft_token
      } else {
        match rng.next_below(config.vocabulary_size - 1) {
          Ok(value) => (draft_token + 1 + value) % config.vocabulary_size
          Err(_) => return Err(RandomFailure)
        }
      }
      draft_logits.push(logits_for(rng, config.vocabulary_size, draft_token))
      target_logits.push(logits_for(rng, config.vocabulary_size, target_token))
    }
    let accept_uniforms = match rng.uniforms(config.proposal_depth) {
      Ok(value) => value
      Err(_) => return Err(RandomFailure)
    }
    let draft_uniforms = match rng.uniforms(config.proposal_depth) {
      Ok(value) => value
      Err(_) => return Err(RandomFailure)
    }
    let fallback_uniforms = match rng.uniforms(config.proposal_depth) {
      Ok(value) => value
      Err(_) => return Err(RandomFailure)
    }
    schedule.push({
      draft_logits,
      target_logits,
      draft_uniforms,
      accept_uniforms,
      fallback_uniforms,
    })
  }
  Ok(schedule)
}

///| Describe a generated workload's structural properties without dumping its

///|
/// potentially large logits. It is suitable for a README or CI log.
pub fn describe_workload(config : WorkloadConfig, seed : Int) -> String {
  "seed=" +
  seed.to_string() +
  " rounds=" +
  config.rounds.to_string() +
  " depth=" +
  config.proposal_depth.to_string() +
  " vocabulary=" +
  config.vocabulary_size.to_string() +
  " agreement=" +
  config.agreement.to_string()
}