frink-server 0.34.0

//! Rendering per-token distributions into OpenAI's `logprobs` shape.
//!
//! The distributions come from the sampler itself
//! (`frink_models::Sampler::sample_reporting`), which returns the very
//! vector a token was drawn from rather than a second opinion computed
//! beside it. This module only shapes them for the wire, and holds the
//! two decisions that shaping forces.
//!
//! # `log(0)` is `-inf`, and `-inf` is not JSON
//!
//! A filtered-out candidate has probability exactly zero -- that is
//! what "the chain removed it" means -- and `ln(0)` is negative
//! infinity, which `serde_json` cannot represent. Such a candidate is
//! OMITTED from `top_logprobs` rather than serialised as `null` or as
//! some large negative stand-in: it was not a candidate, so reporting a
//! number for it would describe a choice the sampler could not have
//! made.
//!
//! The SAMPLED token can never be one of these. It was drawn from this
//! distribution, so its probability is positive by construction, and
//! `token_logprobs` therefore always has a real number in it.
//!
//! # `text_offset` is bytes, and the tokens are the server's own
//!
//! OpenAI's field is a byte offset into the returned text. It is
//! computed here by decoding each token in order and accumulating
//! lengths, rather than by re-tokenizing the finished string: a
//! detokenize-then-retokenize round trip is not the identity for every
//! vocabulary, and an offset that disagrees with the text it indexes is
//! worse than no offset.

use serde_json::Value;

use crate::sampling_loop::PerTokenProbs;

/// The `logprobs` object for one choice, or `Value::Null` when the
/// request did not ask for it.
///
/// `top_k` is the caller's `logprobs: N` -- how many alternatives to
/// report per position. Zero means "the chosen token only", which is
/// what `logprobs: 0` asks for upstream.
pub(crate) fn render(
    per_token: &PerTokenProbs,
    top_k: Option<usize>,
    decode: &dyn Fn(usize) -> String,
) -> Value {
    let Some(top_k) = top_k else {
        return Value::Null;
    };
    if per_token.is_empty() {
        // Asked for, and there is nothing to report: an empty object
        // rather than null, because the request WAS honoured and the
        // generation simply produced no tokens.
        return serde_json::json!({
            "tokens": [],
            "token_logprobs": [],
            "top_logprobs": [],
            "text_offset": [],
        });
    }

    let mut tokens: Vec<String> = Vec::with_capacity(per_token.len());
    let mut token_logprobs: Vec<f64> = Vec::with_capacity(per_token.len());
    let mut top_logprobs: Vec<Value> = Vec::with_capacity(per_token.len());
    let mut text_offset: Vec<usize> = Vec::with_capacity(per_token.len());
    let mut offset = 0usize;

    for (id, probs) in per_token {
        let piece = decode(*id);
        text_offset.push(offset);
        offset += piece.len();
        tokens.push(piece);
        // Positive by construction: this token was drawn from this
        // distribution. `max(f64::MIN)` is not a fallback for that, it
        // guards a denormal that rounds to zero in the f32 -> f64 hop.
        let chosen = probs.get(*id).copied().unwrap_or(0.0);
        token_logprobs.push(logprob(chosen));

        let mut alternatives: Vec<(usize, f32)> = probs
            .iter()
            .enumerate()
            // A zero is a candidate the chain removed, and `ln(0)` is
            // not a number JSON can carry. Omitted, not stand-in'd.
            .filter(|(_, p)| **p > 0.0)
            .map(|(i, p)| (i, *p))
            .collect();
        alternatives.sort_by(|a, b| b.1.total_cmp(&a.1).then(a.0.cmp(&b.0)));
        alternatives.truncate(top_k);
        let mut map = serde_json::Map::new();
        for (alt, p) in alternatives {
            map.insert(decode(alt), Value::from(logprob(p)));
        }
        top_logprobs.push(Value::Object(map));
    }

    serde_json::json!({
        "tokens": tokens,
        "token_logprobs": token_logprobs,
        "top_logprobs": top_logprobs,
        "text_offset": text_offset,
    })
}

/// `ln(p)` in f64, so the value survives JSON.
///
/// A zero cannot reach here from the chosen token, and is filtered out
/// of the alternatives, so this never returns `-inf` for a value that
/// is serialised.
fn logprob(p: f32) -> f64 {
    (p as f64).ln()
}

#[cfg(test)]
mod tests {
    use super::*;

    fn decoder() -> impl Fn(usize) -> String {
        |id: usize| format!("t{id}")
    }

    #[test]
    fn no_request_means_no_object() {
        assert_eq!(render(&Vec::new(), None, &decoder()), Value::Null);
    }

    /// The chosen token's logprob is always a real number, and the
    /// alternatives are sorted by probability with the zeros left out.
    #[test]
    fn the_chosen_token_and_its_alternatives_are_reported() {
        // token 1 chosen from a distribution where token 3 was removed.
        let per_token: PerTokenProbs = vec![(1, vec![0.25, 0.5, 0.25, 0.0])];
        let out = render(&per_token, Some(3), &decoder());

        assert_eq!(out["tokens"], serde_json::json!(["t1"]));
        assert_eq!(out["text_offset"], serde_json::json!([0]));
        let chosen = out["token_logprobs"][0].as_f64().expect("a real number");
        assert!((chosen - 0.5f64.ln()).abs() < 1e-9, "{chosen}");

        let top = out["top_logprobs"][0].as_object().expect("an object");
        assert_eq!(top.len(), 3, "three survived the filter: {top:?}");
        assert!(
            !top.contains_key("t3"),
            "a zero-probability candidate was reported: {top:?}"
        );
        // Sorted by probability, so the most likely is the chosen one.
        assert!((top["t1"].as_f64().unwrap() - 0.5f64.ln()).abs() < 1e-9);
    }

    /// A removed candidate is OMITTED even when there is room for it.
    ///
    /// The earlier version of this test asked for 3 out of 4 where the
    /// zero sorted last, so keeping the zeros would still have
    /// produced 3 entries and the test could not fail. Here `top_k`
    /// exceeds the number of surviving candidates, so the filter is
    /// the only thing that decides the length.
    #[test]
    fn a_removed_candidate_is_omitted_even_with_room_to_spare() {
        let per_token: PerTokenProbs = vec![(0, vec![0.7, 0.3, 0.0, 0.0])];
        let out = render(&per_token, Some(4), &decoder());
        let top = out["top_logprobs"][0].as_object().expect("an object");
        assert_eq!(
            top.len(),
            2,
            "a zero-probability candidate was reported: {top:?}"
        );
        assert!(
            !top.contains_key("t2") && !top.contains_key("t3"),
            "{top:?}"
        );
        for v in top.values() {
            assert!(
                v.as_f64().expect("a real number").is_finite(),
                "ln(0) reached the wire: {top:?}"
            );
        }
    }

    /// `top_k` truncates, and truncates to the MOST likely.
    #[test]
    fn top_k_keeps_the_likeliest() {
        let per_token: PerTokenProbs = vec![(0, vec![0.6, 0.3, 0.1])];
        let out = render(&per_token, Some(2), &decoder());
        let top = out["top_logprobs"][0].as_object().unwrap();
        assert_eq!(top.len(), 2);
        assert!(top.contains_key("t0") && top.contains_key("t1"));
        assert!(!top.contains_key("t2"), "kept the least likely: {top:?}");
    }

    /// Offsets are byte offsets into the concatenated text, computed
    /// from the same pieces the `tokens` array reports -- so a client
    /// can slice the returned text with them.
    #[test]
    fn text_offsets_index_the_text_they_describe() {
        let per_token: PerTokenProbs = vec![
            (1, vec![0.0, 1.0]),
            (0, vec![1.0, 0.0]),
            (1, vec![0.0, 1.0]),
        ];
        let out = render(&per_token, Some(1), &decoder());
        let tokens: Vec<String> = serde_json::from_value(out["tokens"].clone()).unwrap();
        let offsets: Vec<usize> = serde_json::from_value(out["text_offset"].clone()).unwrap();
        let text = tokens.concat();
        for (i, off) in offsets.iter().enumerate() {
            assert!(
                text[*off..].starts_with(&tokens[i]),
                "offset {off} does not point at {:?} in {text:?}",
                tokens[i]
            );
        }
    }

    /// Asked for, nothing generated: the request was honoured, so the
    /// object exists and is empty rather than absent.
    #[test]
    fn an_empty_generation_still_reports_the_object() {
        let out = render(&Vec::new(), Some(2), &decoder());
        assert!(out.is_object(), "{out}");
        assert_eq!(out["tokens"], serde_json::json!([]));
    }
}