「AIに何を入力すればいいか分からない」をなくす技術
指示文不要の業務AIと「制約付きデコーディング」の科学

#![forbid(unsafe_code)]

pub const NEG_INFINITY_LOGIT: f32 = f32::NEG_INFINITY;

#[derive(Debug, Clone)]
pub struct DfaStateMask {
    pub state_id: u32,
    pub is_accepting: bool,
    /// Bitset over vocabulary (1 bit per token: 1 = allowed, 0 = forbidden)
    pub allowed_token_bits: Vec<u64>,
}

impl DfaStateMask {
    pub fn new(state_id: u32, vocab_size: usize, is_accepting: bool) -> Self {
        let words = (vocab_size + 63) / 64;
        Self {
            state_id,
            is_accepting,
            allowed_token_bits: vec![0u64; words],
        }
    }

    #[inline]
    pub fn allow_token(&mut self, token_id: usize) {
        let word = token_id / 64;
        let bit = token_id % 64;
        if word < self.allowed_token_bits.len() {
            self.allowed_token_bits[word] |= 1u64 << bit;
        }
    }

    #[inline]
    pub fn is_token_allowed(&self, token_id: usize) -> bool {
        let word = token_id / 64;
        let bit = token_id % 64;
        self.allowed_token_bits
            .get(word)
            .map_or(false, |w| (w & (1u64 << bit)) != 0)
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MaskError {
    VocabSizeMismatch,
    NoValidTokensInState { state_id: u32 },
}

pub fn apply_dfa_logit_mask(logits: &mut [f32], mask: &DfaStateMask) -> Result<usize, MaskError> {
    let max_tokens = mask.allowed_token_bits.len() * 64;
    if logits.len() > max_tokens {
        return Err(MaskError::VocabSizeMismatch);
    }

    let mut allowed_count = 0usize;
    for (token_id, logit) in logits.iter_mut().enumerate() {
        if mask.is_token_allowed(token_id) {
            allowed_count += 1;
        } else {
            *logit = NEG_INFINITY_LOGIT;
        }
    }

    if allowed_count == 0 {
        return Err(MaskError::NoValidTokensInState {
            state_id: mask.state_id,
        });
    }

    Ok(allowed_count)
}

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

    #[test]
    fn test_forbidden_tokens_are_masked_to_neg_infinity() {
        let mut mask = DfaStateMask::new(1, 128, false);
        mask.allow_token(10);
        mask.allow_token(42);

        let mut logits = vec![1.5f32; 128];
        let allowed = apply_dfa_logit_mask(&mut logits, &mask).unwrap();
        assert_eq!(allowed, 2);
        assert_eq!(logits[10], 1.5f32);
        assert_eq!(logits[42], 1.5f32);
        assert!(logits[0].is_infinite() && logits[0].is_sign_negative());
        assert!(logits[127].is_infinite() && logits[127].is_sign_negative());
    }
}