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.allows(token) {
89            Err(TokenChoiceError::Forbidden(token))
90        } else {
91            Ok(())
92        }
93    }
94}
95impl<C: TokenFilterController> TokenFilterController for TokenChoiceController<C> {
96    type Error = TokenChoiceError<C::Error>;
97    fn current_filter(&mut self) -> Result<TokenFilter, Self::Error> {
98        let filter = self
99            .inner
100            .current_filter()
101            .map_err(TokenChoiceError::Constraint)?;
102        let Some(token) = self.pending else {
103            return Ok(filter);
104        };
105        Self::check_filter(&filter, token)?;
106        let mut allowed = vec![
107            false;
108            filter
109                .allowed_mask()
110                .map_or(self.domain.cardinality(), <[bool]>::len)
111        ];
112        allowed[token as usize] = true;
113        Ok(TokenFilter::allowed(allowed).expect("one canonical candidate is allowed"))
114    }
115    fn commit_token(&mut self, token: u32) -> Result<(), Self::Error> {
116        if let Some(expected) = self.pending {
117            if expected != token {
118                return Err(TokenChoiceError::UnexpectedCommit {
119                    expected,
120                    actual: token,
121                });
122            }
123        }
124        self.inner
125            .commit_token(token)
126            .map_err(TokenChoiceError::Constraint)?;
127        self.last_forced = self.pending.take().is_some();
128        Ok(())
129    }
130    fn is_complete(&mut self) -> Result<bool, Self::Error> {
131        self.inner
132            .is_complete()
133            .map_err(TokenChoiceError::Constraint)
134    }
135}
136impl<C: SnapshotTokenController> SnapshotTokenController for TokenChoiceController<C> {
137    fn snapshot_storage_bytes(&self) -> Option<u64> {
138        self.inner
139            .snapshot_storage_bytes()?
140            .checked_add(std::mem::size_of::<Self>() as u64)
141    }
142    fn fork_snapshot(&self) -> Result<Self, String> {
143        Ok(Self {
144            inner: self.inner.fork_snapshot()?,
145            domain: self.domain,
146            pending: self.pending,
147            last_forced: self.last_forced,
148        })
149    }
150}
151
152#[cfg(test)]
153mod tests {
154    use super::*;
155    use std::convert::Infallible;
156
157    #[derive(Clone)]
158    struct Grammar(Vec<u32>);
159    impl TokenFilterController for Grammar {
160        type Error = Infallible;
161        fn current_filter(&mut self) -> Result<TokenFilter, Infallible> {
162            let mut allowed = vec![true; 4];
163            allowed[self.0.len() % 4] = false;
164            Ok(TokenFilter::allowed(allowed).unwrap())
165        }
166        fn commit_token(&mut self, token: u32) -> Result<(), Infallible> {
167            self.0.push(token);
168            Ok(())
169        }
170        fn is_complete(&mut self) -> Result<bool, Infallible> {
171            Ok(false)
172        }
173    }
174    impl SnapshotTokenController for Grammar {
175        fn snapshot_storage_bytes(&self) -> Option<u64> {
176            Some(24 + 4 * self.0.len() as u64)
177        }
178        fn fork_snapshot(&self) -> Result<Self, String> {
179            Ok(self.clone())
180        }
181    }
182
183    #[test]
184    fn choices_validate_constraints_and_commit_once_with_snapshot_isolation() {
185        let mut parent = TokenChoiceController::new(Grammar(vec![]), TokenDomain::new(4));
186        assert!(matches!(
187            parent.force_next(4),
188            Err(TokenChoiceError::InvalidToken(4))
189        ));
190        assert!(matches!(
191            parent.force_next(0),
192            Err(TokenChoiceError::Forbidden(0))
193        ));
194        assert!(parent.pending_forced().is_none());
195        parent.force_next(2).unwrap();
196        assert_eq!(
197            parent.current_filter().unwrap().allowed_mask(),
198            Some(&[false, false, true, false][..])
199        );
200        assert!(matches!(
201            parent.force_next(3),
202            Err(TokenChoiceError::AlreadyPending)
203        ));
204        let mut child = parent.fork_snapshot().unwrap();
205        assert!(matches!(
206            parent.commit_token(3),
207            Err(TokenChoiceError::UnexpectedCommit { .. })
208        ));
209        assert!(parent.inner().0.is_empty());
210        parent.commit_token(2).unwrap();
211        assert!(parent.last_committed_was_forced());
212        assert!(parent.pending_forced().is_none());
213        assert_eq!(parent.inner().0, [2]);
214        assert_eq!(child.pending_forced(), Some(2));
215        assert!(child.clear_forced());
216        child.commit_token(3).unwrap();
217        assert!(!child.last_committed_was_forced());
218        assert_eq!(child.inner().0, [3]);
219        assert_eq!(parent.inner().0, [2]);
220        assert!(matches!(
221            parent.force_next(1),
222            Err(TokenChoiceError::Forbidden(1))
223        ));
224        parent.commit_token(3).unwrap();
225        assert!(!parent.last_committed_was_forced());
226        assert_eq!(parent.inner().0, [2, 3]);
227    }
228}
eredu_runtime/execution_control/choice.rs

eredu_runtime/execution_control/
choice.rs
