eredu_runtime/execution_control/
choice.rs

1//! One prospective canonical choice through the existing ordinary sampler.
2use super::SnapshotTokenController;
3use crate::TokenDomain;
4use eredu_core::{TokenFilter, TokenFilterController};
5
6/// Rejected canonical choice or error from the original grammar controller.
7#[derive(Debug, thiserror::Error)]
8pub enum TokenChoiceError<E: std::error::Error + 'static> {
9    /// The existing grammar/filter failed.
10    #[error("token constraint failed: {0}")]
11    Constraint(#[source] E),
12    /// ID is outside the canonical vocabulary domain.
13    #[error("forced token {0} is outside the canonical vocabulary")]
14    InvalidToken(u32),
15    /// The active grammar/filter disallows this canonical ID.
16    #[error("forced token {0} conflicts with the active token constraints")]
17    Forbidden(u32),
18    /// A choice is already waiting for commitment.
19    #[error("a forced choice is already pending; clear it before replacing it")]
20    AlreadyPending,
21    /// Backend commitment violated the filter it received.
22    #[error("forced token {expected} was requested, but backend committed {actual}")]
23    UnexpectedCommit {
24        /// The sole allowed canonical ID.
25        expected: u32,
26        /// The ID returned by the backend.
27        actual: u32,
28    },
29}
30
31/// Adds a one-decision restriction to an existing canonical constraint owner.
32/// Native sampling, penalties, history and RNG still use the ordinary path. A
33/// forced decision consumes exactly that sampler's one-candidate decision: at
34/// nonzero temperature RNG advances as usual, and Mirostat observes probability
35/// one. It does not edit an already committed prefix or replay text.
36#[derive(Clone, Debug, PartialEq)]
37pub struct TokenChoiceController<C> {
38    inner: C,
39    domain: TokenDomain,
40    pending: Option<u32>,
41    last_forced: bool,
42}
43
44impl<C: TokenFilterController> TokenChoiceController<C> {
45    /// Wraps the existing grammar with the canonical tokenizer vocabulary domain.
46    pub fn new(inner: C, domain: TokenDomain) -> Self {
47        Self {
48            inner,
49            domain,
50            pending: None,
51            last_forced: false,
52        }
53    }
54    /// Stages a choice after checking vocabulary and the current active grammar.
55    /// No token is committed and no sampler/RNG or model state advances here.
56    pub fn force_next(&mut self, token: u32) -> Result<(), TokenChoiceError<C::Error>> {
57        if self.pending.is_some() {
58            return Err(TokenChoiceError::AlreadyPending);
59        }
60        if token as usize >= self.domain.cardinality() {
61            return Err(TokenChoiceError::InvalidToken(token));
62        }
63        let filter = self
64            .inner
65            .current_filter()
66            .map_err(TokenChoiceError::Constraint)?;
67        Self::check_filter(&filter, token)?;
68        self.pending = Some(token);
69        Ok(())
70    }
71    /// Removes an uncommitted choice without changing history or randomness.
72    pub fn clear_forced(&mut self) -> bool {
73        self.pending.take().is_some()
74    }
75    /// Choice still waiting for the next ordinary commitment.
76    pub fn pending_forced(&self) -> Option<u32> {
77        self.pending
78    }
79    /// Whether the most recent canonical commitment consumed a forced choice.
80    pub fn last_committed_was_forced(&self) -> bool {
81        self.last_forced
82    }
83    /// Read-only access to the original canonical grammar owner.
84    pub fn inner(&self) -> &C {
85        &self.inner
86    }
87    fn check_filter(filter: &TokenFilter, token: u32) -> Result<(), TokenChoiceError<C::Error>> {
88        if filter
89            .allowed_mask()
90            .is_some_and(|mask| !mask.get(token as usize).copied().unwrap_or(false))
91        {
92            Err(TokenChoiceError::Forbidden(token))
93        } else {
94            Ok(())
95        }
96    }
97}
98impl<C: TokenFilterController> TokenFilterController for TokenChoiceController<C> {
99    type Error = TokenChoiceError<C::Error>;
100    fn current_filter(&mut self) -> Result<TokenFilter, Self::Error> {
101        let filter = self
102            .inner
103            .current_filter()
104            .map_err(TokenChoiceError::Constraint)?;
105        let Some(token) = self.pending else {
106            return Ok(filter);
107        };
108        Self::check_filter(&filter, token)?;
109        let mut allowed = vec![
110            false;
111            filter
112                .allowed_mask()
113                .map_or(self.domain.cardinality(), <[bool]>::len)
114        ];
115        allowed[token as usize] = true;
116        Ok(TokenFilter::allowed(allowed).expect("one canonical candidate is allowed"))
117    }
118    fn commit_token(&mut self, token: u32) -> Result<(), Self::Error> {
119        if let Some(expected) = self.pending {
120            if expected != token {
121                return Err(TokenChoiceError::UnexpectedCommit {
122                    expected,
123                    actual: token,
124                });
125            }
126        }
127        self.inner
128            .commit_token(token)
129            .map_err(TokenChoiceError::Constraint)?;
130        self.last_forced = self.pending.take().is_some();
131        Ok(())
132    }
133    fn is_complete(&mut self) -> Result<bool, Self::Error> {
134        self.inner
135            .is_complete()
136            .map_err(TokenChoiceError::Constraint)
137    }
138}
139impl<C: SnapshotTokenController> SnapshotTokenController for TokenChoiceController<C> {
140    fn snapshot_storage_bytes(&self) -> Option<u64> {
141        self.inner
142            .snapshot_storage_bytes()?
143            .checked_add(std::mem::size_of::<Self>() as u64)
144    }
145    fn fork_snapshot(&self) -> Result<Self, String> {
146        Ok(Self {
147            inner: self.inner.fork_snapshot()?,
148            domain: self.domain,
149            pending: self.pending,
150            last_forced: self.last_forced,
151        })
152    }
153}
154
155#[cfg(test)]
156mod tests {
157    use super::*;
158    use std::convert::Infallible;
159
160    #[derive(Clone)]
161    struct Grammar(Vec<u32>);
162    impl TokenFilterController for Grammar {
163        type Error = Infallible;
164        fn current_filter(&mut self) -> Result<TokenFilter, Infallible> {
165            let mut allowed = vec![true; 4];
166            allowed[self.0.len() % 4] = false;
167            Ok(TokenFilter::allowed(allowed).unwrap())
168        }
169        fn commit_token(&mut self, token: u32) -> Result<(), Infallible> {
170            self.0.push(token);
171            Ok(())
172        }
173        fn is_complete(&mut self) -> Result<bool, Infallible> {
174            Ok(false)
175        }
176    }
177    impl SnapshotTokenController for Grammar {
178        fn snapshot_storage_bytes(&self) -> Option<u64> {
179            Some(24 + 4 * self.0.len() as u64)
180        }
181        fn fork_snapshot(&self) -> Result<Self, String> {
182            Ok(self.clone())
183        }
184    }
185
186    #[test]
187    fn choices_validate_constraints_and_commit_once_with_snapshot_isolation() {
188        let mut parent = TokenChoiceController::new(Grammar(vec![]), TokenDomain::new(4));
189        assert!(matches!(
190            parent.force_next(4),
191            Err(TokenChoiceError::InvalidToken(4))
192        ));
193        assert!(matches!(
194            parent.force_next(0),
195            Err(TokenChoiceError::Forbidden(0))
196        ));
197        assert!(parent.pending_forced().is_none());
198        parent.force_next(2).unwrap();
199        assert_eq!(
200            parent.current_filter().unwrap().allowed_mask(),
201            Some(&[false, false, true, false][..])
202        );
203        assert!(matches!(
204            parent.force_next(3),
205            Err(TokenChoiceError::AlreadyPending)
206        ));
207        let mut child = parent.fork_snapshot().unwrap();
208        assert!(matches!(
209            parent.commit_token(3),
210            Err(TokenChoiceError::UnexpectedCommit { .. })
211        ));
212        assert!(parent.inner().0.is_empty());
213        parent.commit_token(2).unwrap();
214        assert!(parent.last_committed_was_forced());
215        assert!(parent.pending_forced().is_none());
216        assert_eq!(parent.inner().0, [2]);
217        assert_eq!(child.pending_forced(), Some(2));
218        assert!(child.clear_forced());
219        child.commit_token(3).unwrap();
220        assert!(!child.last_committed_was_forced());
221        assert_eq!(child.inner().0, [3]);
222        assert_eq!(parent.inner().0, [2]);
223        assert!(matches!(
224            parent.force_next(1),
225            Err(TokenChoiceError::Forbidden(1))
226        ));
227        parent.commit_token(3).unwrap();
228        assert!(!parent.last_committed_was_forced());
229        assert_eq!(parent.inner().0, [2, 3]);
230    }
231}
eredu_runtime/execution_control/choice.rs

eredu_runtime/execution_control/
choice.rs
