eredu_runtime/
generation.rs

1//! Backend-neutral causal-model and token-sampling contracts.
2
3use eredu_core::{
4    generation::ResolvedGenerationConfig, SpeculativeTokenFilterController, TokenFilter,
5    TokenFilterController,
6};
7use eredu_nn::Tensor;
8
9/// Monomorphized causal model used by generation sessions.
10pub trait CausalModel<S> {
11    /// Backend-native tensor handle containing logits and decode token ids.
12    type Tensor: Tensor;
13    /// Borrowed, tokenizer/media-prepared prefill input.
14    type Input<'a>: Copy;
15    /// Concrete model or backend failure.
16    type Error;
17
18    /// Computes initial logits and updates mutable state.
19    fn prefill_input_logits(
20        &mut self,
21        input: Self::Input<'_>,
22        state: &mut S,
23        context: &<Self::Tensor as Tensor>::Context,
24    ) -> Result<Self::Tensor, Self::Error>;
25
26    /// Computes logits for decode tokens using existing mutable state.
27    fn decode_logits(
28        &mut self,
29        input_tokens: &Self::Tensor,
30        state: &mut S,
31        context: &<Self::Tensor as Tensor>::Context,
32    ) -> Result<Self::Tensor, Self::Error>;
33
34    /// Adjusts prefill logits before backend-native sampling.
35    fn adjust_prefill_logits(
36        &mut self,
37        logits: Self::Tensor,
38        _state: &mut S,
39        _context: &<Self::Tensor as Tensor>::Context,
40    ) -> Result<Self::Tensor, Self::Error> {
41        Ok(logits)
42    }
43}
44
45/// Portable cardinality of one zero-based token-id domain.
46///
47/// Backends validate native token tensors without copying their values to the
48/// host. Architectures select the exact domain for each prediction boundary.
49#[derive(Debug, Clone, Copy, Eq, PartialEq)]
50pub struct TokenDomain {
51    cardinality: usize,
52}
53
54impl TokenDomain {
55    /// Creates the token IDs `0..cardinality`.
56    pub const fn new(cardinality: usize) -> Self {
57        Self { cardinality }
58    }
59
60    /// Number of valid zero-based token IDs.
61    pub const fn cardinality(self) -> usize {
62        self.cardinality
63    }
64}
65
66/// Backend primitives required by generic token-sampling policies.
67///
68/// The runtime owns ordering, history, adaptive state, and constraint rollback.
69/// Implementations operate directly on native logits and random state without
70/// copying values through a neutral tensor representation.
71pub trait SamplingBackend {
72    /// Backend-native logits tensor.
73    type Logits: Clone;
74    /// Backend-native sampled-token tensor.
75    type Token: Clone;
76    /// Backend-native random-key stream.
77    type RandomState;
78    /// Execution context, such as a stream.
79    type Context: ?Sized;
80    /// Backend failure.
81    type Error;
82
83    /// Creates a backend error for a portable policy or constraint failure.
84    fn error(message: String) -> Self::Error;
85
86    /// Validates a native token tensor against one architecture-selected domain.
87    ///
88    /// The returned token must retain a backend-native dependency on the range
89    /// check so lazy backends cannot commit an unchecked forced token.
90    fn validate_token(
91        token: &Self::Token,
92        domain: TokenDomain,
93        context: &Self::Context,
94    ) -> Result<Self::Token, Self::Error>;
95
96    /// Scales logits by inverse temperature, preserving the native tensor.
97    fn scale_temperature(
98        logits: &Self::Logits,
99        temperature: f32,
100        context: &Self::Context,
101    ) -> Result<Self::Logits, Self::Error>;
102
103    /// Applies repetition, frequency, and presence penalties.
104    fn apply_penalties(
105        logits: &Self::Logits,
106        history: &[u32],
107        penalties: PenaltyConfig,
108        context: &Self::Context,
109    ) -> Result<Self::Logits, Self::Error>;
110
111    /// Masks all but the highest `top_k` logits. Non-positive values disable it.
112    fn apply_top_k(
113        logits: Self::Logits,
114        top_k: i32,
115        context: &Self::Context,
116    ) -> Result<Self::Logits, Self::Error>;
117
118    /// Applies nucleus filtering while retaining canonical vocabulary order.
119    fn apply_top_p(
120        logits: Self::Logits,
121        top_p: f32,
122        context: &Self::Context,
123    ) -> Result<Self::Logits, Self::Error>;
124
125    /// Applies minimum-relative-probability filtering.
126    fn apply_min_p(
127        logits: Self::Logits,
128        min_p: f32,
129        context: &Self::Context,
130    ) -> Result<Self::Logits, Self::Error>;
131
132    /// Masks tokens rejected by a portable vocabulary filter.
133    /// Explicit masks are closed sets: wider logits must mask the missing IDs,
134    /// while narrower logits use the executable prefix. Reject an empty
135    /// intersection before sampling; [`TokenFilter::allowed_mask_for`] implements
136    /// this shared domain policy without backend tensors.
137    fn apply_token_filter(
138        logits: &Self::Logits,
139        filter: &TokenFilter,
140        context: &Self::Context,
141    ) -> Result<Self::Logits, Self::Error>;
142
143    /// Applies Mirostat's surprise cutoff after penalties and temperature.
144    fn apply_mirostat(
145        logits: &Self::Logits,
146        history: &[u32],
147        penalties: PenaltyConfig,
148        temperature: f32,
149        mu: f32,
150        context: &Self::Context,
151    ) -> Result<Self::Logits, Self::Error>;
152
153    /// Selects from raw logits, applying temperature for stochastic sampling.
154    fn sample_raw(
155        logits: &Self::Logits,
156        temperature: f32,
157        random: Option<&mut Self::RandomState>,
158        context: &Self::Context,
159    ) -> Result<Self::Token, Self::Error>;
160
161    /// Selects from logits already scaled by the policy.
162    fn sample_processed(
163        logits: &Self::Logits,
164        temperature: f32,
165        random: Option<&mut Self::RandomState>,
166        context: &Self::Context,
167    ) -> Result<Self::Token, Self::Error>;
168
169    /// Materializes only the selected scalar token identifier.
170    fn token_id(token: &Self::Token, context: &Self::Context) -> Result<u32, Self::Error>;
171
172    /// Materializes one committed token probability from processed logits.
173    fn token_probability(
174        logits: &Self::Logits,
175        token: u32,
176        context: &Self::Context,
177    ) -> Result<f32, Self::Error>;
178}
179
180/// Backend-neutral repetition/frequency/presence controls.
181#[derive(Debug, Clone, Copy, PartialEq)]
182pub struct PenaltyConfig {
183    /// Repetition multiplier; `1.0` disables it.
184    pub repeat_penalty: f32,
185    /// Number of recent tokens considered; negative means all.
186    pub repeat_last_n: i32,
187    /// Per-occurrence logit penalty.
188    pub frequency_penalty: f32,
189    /// One-time penalty for every present token.
190    pub presence_penalty: f32,
191}
192
193impl PenaltyConfig {
194    /// Returns whether every penalty is disabled.
195    pub fn is_identity(self) -> bool {
196        self.repeat_penalty == 1.0 && self.frequency_penalty == 0.0 && self.presence_penalty == 0.0
197    }
198}
199
200impl Default for PenaltyConfig {
201    fn default() -> Self {
202        Self {
203            repeat_penalty: 1.0,
204            repeat_last_n: 64,
205            frequency_penalty: 0.0,
206            presence_penalty: 0.0,
207        }
208    }
209}
210
211/// Sampling policy suitable for lossless speculative decoding.
212pub trait SpeculativeSampler<B: SamplingBackend> {
213    /// Whether loaded checkpoint defaults should wrap this policy.
214    fn uses_checkpoint_defaults(&self) -> bool {
215        false
216    }
217
218    /// Whether optimistic draft work is an exact discardable fork.
219    fn supports_exact_optimistic_promotion(&self) -> bool {
220        false
221    }
222
223    /// Whether the committed generation grammar is complete.
224    fn grammar_is_complete(&mut self) -> Result<bool, B::Error> {
225        Ok(false)
226    }
227
228    /// Whether an uncommitted logical prefix completes the grammar.
229    fn prefix_is_complete(&self, _history: &[u32]) -> Result<bool, B::Error> {
230        Ok(false)
231    }
232
233    /// Applies penalties, filters, and temperature.
234    fn process_logits(
235        &mut self,
236        logits: &B::Logits,
237        temperature: f32,
238        history: &[u32],
239        context: &B::Context,
240    ) -> Result<B::Logits, B::Error>;
241
242    /// Selects from already processed logits.
243    fn sample_processed(
244        &self,
245        logits: &B::Logits,
246        temperature: f32,
247        random: Option<&mut B::RandomState>,
248        context: &B::Context,
249    ) -> Result<B::Token, B::Error> {
250        B::sample_processed(logits, temperature, random, context)
251    }
252
253    /// Commits a token selected from a processed target distribution.
254    fn commit_token(
255        &mut self,
256        _processed_logits: &B::Logits,
257        _token: u32,
258        _context: &B::Context,
259    ) -> Result<(), B::Error> {
260        Ok(())
261    }
262}
263
264/// Strategy for choosing a token from model logits.
265pub trait Sampler<B: SamplingBackend> {
266    /// Whether loaded checkpoint defaults should wrap this policy.
267    fn uses_checkpoint_defaults(&self) -> bool {
268        false
269    }
270
271    /// Selects one token from raw model logits.
272    fn sample(
273        &mut self,
274        logits: &B::Logits,
275        temperature: f32,
276        random: Option<&mut B::RandomState>,
277        context: &B::Context,
278    ) -> Result<B::Token, B::Error>;
279}
280
281/// Grammar-aware wrapper around a backend-neutral sampling policy.
282pub struct ConstrainedSampler<S, C> {
283    policy: S,
284    controller: C,
285}
286
287struct ConstraintCheckpoint<S, C> {
288    policy: S,
289    controller: C,
290}
291
292impl<S: Clone, C: Clone> Clone for ConstrainedSampler<S, C> {
293    fn clone(&self) -> Self {
294        Self {
295            policy: self.policy.clone(),
296            controller: self.controller.clone(),
297        }
298    }
299}
300
301impl<S, C> ConstrainedSampler<S, C> {
302    /// Wraps a policy with a portable canonical constraint controller.
303    pub fn new(policy: S, controller: C) -> Self {
304        Self { policy, controller }
305    }
306
307    /// Returns the wrapped policy.
308    pub const fn policy(&self) -> &S {
309        &self.policy
310    }
311
312    /// Returns the portable constraint controller.
313    pub const fn controller(&self) -> &C {
314        &self.controller
315    }
316
317    /// Returns the portable constraint controller mutably.
318    pub fn controller_mut(&mut self) -> &mut C {
319        &mut self.controller
320    }
321}
322
323impl<S: Clone, C: Clone> ConstrainedSampler<S, C> {
324    fn checkpoint(&self) -> ConstraintCheckpoint<S, C> {
325        ConstraintCheckpoint {
326            policy: self.policy.clone(),
327            controller: self.controller.clone(),
328        }
329    }
330}
331
332impl<B, S, C> SpeculativeSampler<B> for ConstrainedSampler<S, C>
333where
334    B: SamplingBackend,
335    S: SpeculativeSampler<B> + Clone,
336    C: SpeculativeTokenFilterController,
337{
338    fn supports_exact_optimistic_promotion(&self) -> bool {
339        self.policy.supports_exact_optimistic_promotion()
340    }
341
342    fn grammar_is_complete(&mut self) -> Result<bool, B::Error> {
343        self.controller
344            .is_complete()
345            .map_err(|error| B::error(error.to_string()))
346    }
347
348    fn prefix_is_complete(&self, history: &[u32]) -> Result<bool, B::Error> {
349        self.controller
350            .prefix_is_complete(history)
351            .map_err(|error| B::error(error.to_string()))
352    }
353
354    fn process_logits(
355        &mut self,
356        logits: &B::Logits,
357        temperature: f32,
358        history: &[u32],
359        context: &B::Context,
360    ) -> Result<B::Logits, B::Error> {
361        let filter = self
362            .controller
363            .filter_at(history)
364            .map_err(|error| B::error(error.to_string()))?;
365        let masked = B::apply_token_filter(logits, &filter, context)?;
366        self.policy
367            .process_logits(&masked, temperature, history, context)
368    }
369
370    fn sample_processed(
371        &self,
372        logits: &B::Logits,
373        temperature: f32,
374        random: Option<&mut B::RandomState>,
375        context: &B::Context,
376    ) -> Result<B::Token, B::Error> {
377        self.policy
378            .sample_processed(logits, temperature, random, context)
379    }
380
381    fn commit_token(
382        &mut self,
383        processed_logits: &B::Logits,
384        token: u32,
385        context: &B::Context,
386    ) -> Result<(), B::Error> {
387        let checkpoint = self.checkpoint();
388        if let Err(error) = self
389            .policy
390            .commit_token(processed_logits, token, context)
391            .and_then(|()| {
392                self.controller
393                    .commit_token(token)
394                    .map_err(|error| B::error(error.to_string()))
395            })
396        {
397            self.policy = checkpoint.policy;
398            self.controller = checkpoint.controller;
399            return Err(error);
400        }
401        Ok(())
402    }
403}
404
405impl<B, S, C> Sampler<B> for ConstrainedSampler<S, C>
406where
407    B: SamplingBackend,
408    S: Sampler<B> + Clone,
409    C: TokenFilterController + Clone,
410{
411    fn sample(
412        &mut self,
413        logits: &B::Logits,
414        temperature: f32,
415        random: Option<&mut B::RandomState>,
416        context: &B::Context,
417    ) -> Result<B::Token, B::Error> {
418        let checkpoint = self.checkpoint();
419        let filter = self
420            .controller
421            .current_filter()
422            .map_err(|error| B::error(error.to_string()))?;
423        let masked = B::apply_token_filter(logits, &filter, context)?;
424        let token = self.policy.sample(&masked, temperature, random, context)?;
425        let token_id = B::token_id(&token, context)?;
426        if let Err(error) = self.controller.commit_token(token_id) {
427            self.policy = checkpoint.policy;
428            self.controller = checkpoint.controller;
429            return Err(B::error(error.to_string()));
430        }
431        Ok(token)
432    }
433}
434
435/// Stateless greedy/categorical sampler.
436#[derive(Debug, Clone, Copy)]
437pub struct DefaultSampler;
438
439impl<B: SamplingBackend> SpeculativeSampler<B> for DefaultSampler {
440    fn uses_checkpoint_defaults(&self) -> bool {
441        true
442    }
443
444    fn supports_exact_optimistic_promotion(&self) -> bool {
445        true
446    }
447
448    fn process_logits(
449        &mut self,
450        logits: &B::Logits,
451        temperature: f32,
452        _history: &[u32],
453        context: &B::Context,
454    ) -> Result<B::Logits, B::Error> {
455        if temperature == 0.0 {
456            Ok(logits.clone())
457        } else {
458            B::scale_temperature(logits, temperature, context)
459        }
460    }
461}
462
463impl<B: SamplingBackend> Sampler<B> for DefaultSampler {
464    fn uses_checkpoint_defaults(&self) -> bool {
465        true
466    }
467
468    fn sample(
469        &mut self,
470        logits: &B::Logits,
471        temperature: f32,
472        random: Option<&mut B::RandomState>,
473        context: &B::Context,
474    ) -> Result<B::Token, B::Error> {
475        B::sample_raw(logits, temperature, random, context)
476    }
477}
478
479/// Configurable backend-neutral text sampler.
480#[derive(Debug, Clone)]
481pub struct GenerationSampler {
482    /// Keep only the `top_k` highest-logit tokens when positive.
483    pub top_k: i32,
484    /// Nucleus probability mass.
485    pub top_p: f32,
486    /// Minimum probability relative to the most probable token.
487    pub min_p: f32,
488    /// Repetition multiplier.
489    pub repeat_penalty: f32,
490    /// Number of recent tokens considered by penalties.
491    pub repeat_last_n: i32,
492    /// Per-occurrence penalty.
493    pub frequency_penalty: f32,
494    /// One-time presence penalty.
495    pub presence_penalty: f32,
496    generated_tokens: Vec<u32>,
497}
498
499impl Default for GenerationSampler {
500    fn default() -> Self {
501        Self {
502            top_k: 40,
503            top_p: 0.95,
504            min_p: 0.05,
505            repeat_penalty: 1.0,
506            repeat_last_n: 64,
507            frequency_penalty: 0.0,
508            presence_penalty: 0.0,
509            generated_tokens: Vec::new(),
510        }
511    }
512}
513
514impl GenerationSampler {
515    /// Creates a sampler with default controls.
516    pub fn new() -> Self {
517        Self::default()
518    }
519
520    /// Creates a sampler from a resolved portable generation configuration.
521    pub fn from_resolved(config: ResolvedGenerationConfig) -> Self {
522        Self::new()
523            .top_k(config.top_k)
524            .top_p(config.top_p)
525            .min_p(config.min_p)
526            .penalties(
527                config.repetition_penalty,
528                config.repeat_last_n,
529                config.frequency_penalty,
530                config.presence_penalty,
531            )
532    }
533
534    /// Seeds accepted-token history.
535    pub fn with_generated_tokens(mut self, tokens: impl IntoIterator<Item = u32>) -> Self {
536        self.generated_tokens = tokens.into_iter().collect();
537        self
538    }
539
540    /// Sets top-k filtering.
541    pub fn top_k(mut self, value: i32) -> Self {
542        self.top_k = value;
543        self
544    }
545
546    /// Sets nucleus filtering.
547    pub fn top_p(mut self, value: f32) -> Self {
548        self.top_p = value;
549        self
550    }
551
552    /// Sets minimum-relative-probability filtering.
553    pub fn min_p(mut self, value: f32) -> Self {
554        self.min_p = value;
555        self
556    }
557
558    /// Sets repetition, frequency, and presence penalties.
559    pub fn penalties(
560        mut self,
561        repeat_penalty: f32,
562        repeat_last_n: i32,
563        frequency_penalty: f32,
564        presence_penalty: f32,
565    ) -> Self {
566        self.repeat_penalty = repeat_penalty;
567        self.repeat_last_n = repeat_last_n;
568        self.frequency_penalty = frequency_penalty;
569        self.presence_penalty = presence_penalty;
570        self
571    }
572
573    /// Returns accepted-token history.
574    pub fn generated_tokens(&self) -> &[u32] {
575        &self.generated_tokens
576    }
577
578    /// Replaces accepted-token history.
579    pub fn set_generated_tokens(&mut self, tokens: impl IntoIterator<Item = u32>) {
580        self.generated_tokens = tokens.into_iter().collect();
581    }
582
583    /// Records a token accepted outside this sampler.
584    pub fn accept_token(&mut self, token: u32) {
585        self.generated_tokens.push(token);
586    }
587
588    /// Clears accepted-token history.
589    pub fn clear_generated_tokens(&mut self) {
590        self.generated_tokens.clear();
591    }
592
593    /// Returns the portable penalty controls.
594    pub const fn penalty_config(&self) -> PenaltyConfig {
595        PenaltyConfig {
596            repeat_penalty: self.repeat_penalty,
597            repeat_last_n: self.repeat_last_n,
598            frequency_penalty: self.frequency_penalty,
599            presence_penalty: self.presence_penalty,
600        }
601    }
602
603    fn process_for<B: SamplingBackend>(
604        &self,
605        logits: &B::Logits,
606        history: &[u32],
607        context: &B::Context,
608    ) -> Result<B::Logits, B::Error> {
609        let logits = B::apply_penalties(logits, history, self.penalty_config(), context)?;
610        let logits = B::apply_top_k(logits, self.top_k, context)?;
611        let logits = B::apply_top_p(logits, self.top_p, context)?;
612        B::apply_min_p(logits, self.min_p, context)
613    }
614}
615
616impl<B: SamplingBackend> SpeculativeSampler<B> for GenerationSampler {
617    fn supports_exact_optimistic_promotion(&self) -> bool {
618        true
619    }
620
621    fn process_logits(
622        &mut self,
623        logits: &B::Logits,
624        temperature: f32,
625        history: &[u32],
626        context: &B::Context,
627    ) -> Result<B::Logits, B::Error> {
628        let logits = self.process_for::<B>(logits, history, context)?;
629        if temperature == 0.0 {
630            Ok(logits)
631        } else {
632            B::scale_temperature(&logits, temperature, context)
633        }
634    }
635}
636
637impl<B: SamplingBackend> Sampler<B> for GenerationSampler {
638    fn sample(
639        &mut self,
640        logits: &B::Logits,
641        temperature: f32,
642        random: Option<&mut B::RandomState>,
643        context: &B::Context,
644    ) -> Result<B::Token, B::Error> {
645        let logits = self.process_for::<B>(logits, &self.generated_tokens, context)?;
646        let token = B::sample_raw(&logits, temperature, random, context)?;
647        self.generated_tokens.push(B::token_id(&token, context)?);
648        Ok(token)
649    }
650}
651
652/// Adaptive Mirostat V2 policy with backend-neutral state.
653#[derive(Debug, Clone)]
654pub struct MirostatV2Sampler {
655    tau: f32,
656    eta: f32,
657    mu: f32,
658    penalties: GenerationSampler,
659}
660
661impl Default for MirostatV2Sampler {
662    fn default() -> Self {
663        Self {
664            tau: 5.0,
665            eta: 0.1,
666            mu: 10.0,
667            penalties: GenerationSampler::new().top_k(0).top_p(1.0).min_p(0.0),
668        }
669    }
670}
671
672impl MirostatV2Sampler {
673    /// Creates a sampler targeting `tau` bits of surprise.
674    pub fn new(tau: f32, eta: f32) -> Result<Self, SamplingConfigurationError> {
675        validate_positive_finite("Mirostat V2 tau", tau)?;
676        validate_positive_finite("Mirostat V2 eta", eta)?;
677        Ok(Self {
678            tau,
679            eta,
680            mu: 2.0 * tau,
681            penalties: GenerationSampler::new().top_k(0).top_p(1.0).min_p(0.0),
682        })
683    }
684
685    /// Sets penalties applied before adaptive truncation.
686    pub fn penalties(
687        mut self,
688        repeat_penalty: f32,
689        repeat_last_n: i32,
690        frequency_penalty: f32,
691        presence_penalty: f32,
692    ) -> Self {
693        self.penalties = self.penalties.penalties(
694            repeat_penalty,
695            repeat_last_n,
696            frequency_penalty,
697            presence_penalty,
698        );
699        self
700    }
701
702    /// Target surprise in bits.
703    pub const fn tau(&self) -> f32 {
704        self.tau
705    }
706
707    /// Adaptation rate.
708    pub const fn eta(&self) -> f32 {
709        self.eta
710    }
711
712    /// Current adaptive surprise limit.
713    pub const fn mu(&self) -> f32 {
714        self.mu
715    }
716
717    /// Accepted-token history.
718    pub fn generated_tokens(&self) -> &[u32] {
719        self.penalties.generated_tokens()
720    }
721
722    /// Records an externally accepted token and its normalized probability.
723    pub fn accept_token(
724        &mut self,
725        token: u32,
726        probability: f32,
727    ) -> Result<(), SamplingConfigurationError> {
728        if !probability.is_finite() || probability <= 0.0 || probability > 1.0 {
729            return Err(SamplingConfigurationError::Invalid(
730                "accepted Mirostat V2 token probability must be finite and in (0, 1]".into(),
731            ));
732        }
733        self.update_mu(-probability.log2());
734        self.penalties.accept_token(token);
735        Ok(())
736    }
737
738    /// Resets adaptive state and history.
739    pub fn reset(&mut self) {
740        self.mu = 2.0 * self.tau;
741        self.penalties.clear_generated_tokens();
742    }
743
744    fn update_mu(&mut self, observed_surprise: f32) {
745        self.mu -= self.eta * (observed_surprise - self.tau);
746    }
747
748    fn process_for<B: SamplingBackend>(
749        &self,
750        logits: &B::Logits,
751        temperature: f32,
752        history: &[u32],
753        context: &B::Context,
754    ) -> Result<B::Logits, B::Error> {
755        if !temperature.is_finite() || temperature <= 0.0 {
756            return Err(B::error(
757                "Mirostat V2 requires a finite temperature greater than zero".into(),
758            ));
759        }
760        B::apply_mirostat(
761            logits,
762            history,
763            self.penalties.penalty_config(),
764            temperature,
765            self.mu,
766            context,
767        )
768    }
769
770    fn commit_for<B: SamplingBackend>(
771        &mut self,
772        logits: &B::Logits,
773        token: u32,
774        context: &B::Context,
775    ) -> Result<(), B::Error> {
776        let probability = B::token_probability(logits, token, context)?;
777        self.accept_token(token, probability)
778            .map_err(|error| B::error(error.to_string()))
779    }
780}
781
782impl<B: SamplingBackend> Sampler<B> for MirostatV2Sampler {
783    fn sample(
784        &mut self,
785        logits: &B::Logits,
786        temperature: f32,
787        random: Option<&mut B::RandomState>,
788        context: &B::Context,
789    ) -> Result<B::Token, B::Error> {
790        let processed = self.process_for::<B>(
791            logits,
792            temperature,
793            self.penalties.generated_tokens(),
794            context,
795        )?;
796        let token = B::sample_processed(&processed, temperature, random, context)?;
797        self.commit_for::<B>(&processed, B::token_id(&token, context)?, context)?;
798        Ok(token)
799    }
800}
801
802impl<B: SamplingBackend> SpeculativeSampler<B> for MirostatV2Sampler {
803    fn process_logits(
804        &mut self,
805        logits: &B::Logits,
806        temperature: f32,
807        history: &[u32],
808        context: &B::Context,
809    ) -> Result<B::Logits, B::Error> {
810        self.process_for::<B>(logits, temperature, history, context)
811    }
812
813    fn commit_token(
814        &mut self,
815        processed_logits: &B::Logits,
816        token: u32,
817        context: &B::Context,
818    ) -> Result<(), B::Error> {
819        self.commit_for::<B>(processed_logits, token, context)
820    }
821}
822
823/// Invalid backend-neutral sampling configuration.
824#[derive(Debug, Clone, Eq, PartialEq, thiserror::Error)]
825pub enum SamplingConfigurationError {
826    /// A numeric or probability control is invalid.
827    #[error("{0}")]
828    Invalid(String),
829}
830
831fn validate_positive_finite(name: &str, value: f32) -> Result<(), SamplingConfigurationError> {
832    if value.is_finite() && value > 0.0 {
833        Ok(())
834    } else {
835        Err(SamplingConfigurationError::Invalid(format!(
836            "{name} must be finite and greater than zero"
837        )))
838    }
839}
840
841#[cfg(test)]
842mod tests {
843    use super::{GenerationSampler, MirostatV2Sampler};
844
845    #[test]
846    fn generation_history_is_backend_neutral() {
847        let mut sampler = GenerationSampler::new().with_generated_tokens([1, 2]);
848        sampler.accept_token(3);
849        assert_eq!(sampler.generated_tokens(), &[1, 2, 3]);
850        sampler.set_generated_tokens([5, 8]);
851        assert_eq!(sampler.generated_tokens(), &[5, 8]);
852        sampler.clear_generated_tokens();
853        assert!(sampler.generated_tokens().is_empty());
854    }
855
856    #[test]
857    fn mirostat_state_is_backend_neutral() {
858        let mut sampler = MirostatV2Sampler::default();
859        sampler.accept_token(42, 2.0f32.powi(-7)).unwrap();
860        assert!((sampler.mu() - 9.8).abs() < 1e-6);
861        assert_eq!(sampler.generated_tokens(), &[42]);
862        sampler.reset();
863        assert_eq!(sampler.mu(), 10.0);
864        assert!(sampler.generated_tokens().is_empty());
865    }
866}
eredu_runtime/generation.rs

eredu_runtime/
generation.rs
