laya_core/

prompt.rs

//! Prompt construction.
//!
//! Layout, matching upstream exactly:
//!
//! ```text
//! [CLS] <type> question: <instructions> [SEP] [MASK] opt0 [MASK] opt1 ... [SEP] <state> [SEP]
//! ```
//!
//! The `[MASK]` before each option is the marker whose encoder output the
//! scoring head reads, so marker positions and truncation have to agree with
//! upstream token for token.

use crate::error::Result;
use crate::json_compat::dumps_unicode;
use crate::question::Question;
use crate::tokenizer::Tokenizer;
use serde_json::Value;

/// Longest option, in tokens, before truncation.
const MAX_OPTION_TOKENS: usize = 48;
/// Minimum token budget the instructions keep when options are long.
const MIN_HEAD_TOKENS: usize = 8;
/// Below this remaining budget, options are trimmed to make room.
const MIN_OPTION_BUDGET: usize = 16;

/// A tokenized question prefix: ids plus the marker offset of each option.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Prefix {
    pub ids: Vec<u32>,
    pub markers: Vec<usize>,
}

/// A fully built sequence ready for collation.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Sequence {
    pub ids: Vec<u32>,
    pub markers: Vec<usize>,
    /// True when the state did not fit in `max_len` and its tail was dropped.
    ///
    /// Truncation is silent otherwise: the answers stay well formed and the
    /// token count looks ordinary, so a caller that must not send a cut state
    /// has nothing else to check.
    pub truncated: bool,
}

/// Upstream's `serialize_state`: strings pass through, everything else becomes
/// `json.dumps(..., ensure_ascii=False)`.
pub fn serialize_state(state: &Value) -> String {
    match state {
        Value::String(text) => text.clone(),
        other => dumps_unicode(other),
    }
}

/// Build the question-only prefix, before state tokens and final truncation.
pub fn build_prefix(
    tokenizer: &Tokenizer,
    question: &Question,
    head_max_len: usize,
) -> Result<Prefix> {
    let mask = tokenizer.mask_token.as_str();
    let options = question.render_options();

    // Literal mask tokens in user text would create phantom markers.
    let instructions = question.instructions.replace(mask, " ");
    let head_text = format!("{} question: {}", question.kind.name(), instructions);
    let mut head_ids = tokenizer.encode(&head_text)?;

    let mut option_ids: Vec<Vec<u32>> = Vec::with_capacity(options.len());
    for option in &options {
        let text = format!(" {}", option.replace(mask, " "));
        let mut ids = tokenizer.encode(&text)?;
        ids.truncate(MAX_OPTION_TOKENS);
        let mut with_marker = Vec::with_capacity(ids.len() + 1);
        with_marker.push(tokenizer.mask_token_id);
        with_marker.extend_from_slice(&ids);
        option_ids.push(with_marker);
    }

    let used: usize = option_ids.iter().map(Vec::len).sum();
    // `head_max_len - used` is saturating here; upstream lets it go negative and
    // then clamps with `max(8, ...)`, which lands on the same slice length.
    let mut option_budget = head_max_len.saturating_sub(used);
    if option_budget < MIN_OPTION_BUDGET {
        let per = std::cmp::max(
            4,
            head_max_len.saturating_sub(MIN_OPTION_BUDGET) / std::cmp::max(1, option_ids.len()),
        );
        for ids in &mut option_ids {
            ids.truncate(per);
        }
        let used: usize = option_ids.iter().map(Vec::len).sum();
        option_budget = head_max_len.saturating_sub(used);
    }
    head_ids.truncate(std::cmp::max(MIN_HEAD_TOKENS, option_budget));

    let mut ids = Vec::with_capacity(head_ids.len() + used + 3);
    ids.push(tokenizer.cls_token_id);
    ids.extend_from_slice(&head_ids);
    ids.push(tokenizer.sep_token_id);

    let mut markers = Vec::with_capacity(option_ids.len());
    for option in &option_ids {
        markers.push(ids.len());
        ids.extend_from_slice(option);
    }
    ids.push(tokenizer.sep_token_id);

    Ok(Prefix { ids, markers })
}

/// Append state tokens to a prefix and apply the final length cap.
pub fn extend_with_state(
    tokenizer: &Tokenizer,
    prefix: &Prefix,
    state_ids: &[u32],
    max_len: usize,
    truncate_left: bool,
) -> Sequence {
    let room = max_len.saturating_sub(prefix.ids.len() + 1);
    let state = if truncate_left {
        // Upstream writes `st[-room:]`, and in Python `st[-0:]` is the whole
        // list rather than an empty one. Reproduced deliberately; `truncate_left`
        // is not reachable from the public API, so this only matters for parity.
        if room == 0 {
            state_ids
        } else {
            &state_ids[state_ids.len().saturating_sub(room)..]
        }
    } else {
        &state_ids[..std::cmp::min(room, state_ids.len())]
    };

    let mut ids = Vec::with_capacity(prefix.ids.len() + state.len() + 1);
    ids.extend_from_slice(&prefix.ids);
    ids.extend_from_slice(state);
    ids.push(tokenizer.sep_token_id);

    // Either the state was clipped to fit the budget, or the prefix alone
    // already overran it and the final truncate is about to drop real tokens.
    let truncated = state.len() < state_ids.len() || ids.len() > max_len;
    ids.truncate(max_len);

    let markers = prefix
        .markers
        .iter()
        .copied()
        .filter(|marker| *marker < max_len)
        .collect();

    Sequence {
        ids,
        markers,
        truncated,
    }
}

/// Build one complete sequence for a question and state.
pub fn build_sequence(
    tokenizer: &Tokenizer,
    state: &Value,
    question: &Question,
    max_len: usize,
    head_max_len: usize,
) -> Result<Sequence> {
    let prefix = build_prefix(tokenizer, question, head_max_len)?;
    let state_text = serialize_state(state).replace(tokenizer.mask_token.as_str(), " ");
    let state_ids = tokenizer.encode(&state_text)?;
    Ok(extend_with_state(
        tokenizer, &prefix, &state_ids, max_len, false,
    ))
}