Skip to main content

multiscreen_rs/
config.rs

1use crate::error::{Error, Result};
2use crate::screen::ScreenConfig;
3use crate::tile::TileConfig;
4use serde::{Deserialize, Serialize};
5
6/// Multi-screen scoring parameters.
7///
8/// `acceptance_sharpness_r` matches the trim-and-square gate from the experiment:
9/// `square(clamp(1 - acceptance_sharpness_r * (1 - similarity), 0, 1))`.
10///
11/// `screening_window_w` controls the causal softmask width in token positions.
12#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
13pub struct TrimConfig {
14    /// Paper symbol `r`; larger values narrow the acceptance width `1 / r`.
15    #[serde(alias = "radius")]
16    pub acceptance_sharpness_r: f32,
17    /// Paper symbol `w`; causal distance window used by the softmask.
18    #[serde(alias = "window")]
19    pub screening_window_w: f32,
20}
21
22impl Default for TrimConfig {
23    fn default() -> Self {
24        Self {
25            acceptance_sharpness_r: 2.0,
26            screening_window_w: 32.0,
27        }
28    }
29}
30
31impl TrimConfig {
32    pub fn validate(&self) -> Result<()> {
33        if !self.acceptance_sharpness_r.is_finite() || self.acceptance_sharpness_r <= 0.0 {
34            return Err(Error::Config(
35                "trim acceptance_sharpness_r must be positive and finite".into(),
36            ));
37        }
38        if !self.screening_window_w.is_finite() || self.screening_window_w <= 0.0 {
39            return Err(Error::Config(
40                "trim screening_window_w must be positive and finite".into(),
41            ));
42        }
43        Ok(())
44    }
45}
46
47/// Controls deterministic inference behavior for the lightweight transition
48/// engine (`MultiscreenEngine`).
49///
50/// This is distinct from [`ModelInferenceConfig`](crate::ModelInferenceConfig)
51/// which controls the neural `MultiscreenModel` inference path.
52#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
53pub struct InferenceConfig {
54    /// Maximum number of output tokens to produce. `None` means output length
55    /// equals input length.
56    #[serde(alias = "max_output_tokens")]
57    pub max_inference_tokens: Option<usize>,
58    /// When `true` and no trained transition exists for a token, fall back to
59    /// echoing the input token.
60    #[serde(alias = "fallback_to_input")]
61    pub use_input_token_fallback: bool,
62}
63
64impl Default for InferenceConfig {
65    fn default() -> Self {
66        Self {
67            max_inference_tokens: None,
68            use_input_token_fallback: true,
69        }
70    }
71}
72
73impl InferenceConfig {
74    pub fn validate(&self) -> Result<()> {
75        if matches!(self.max_inference_tokens, Some(0)) {
76            return Err(Error::Config(
77                "max_inference_tokens must be omitted or greater than zero".into(),
78            ));
79        }
80        Ok(())
81    }
82}
83
84/// Full engine configuration.
85#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
86pub struct MultiscreenConfig {
87    pub screens: ScreenConfig,
88    pub tiles: TileConfig,
89    pub trim: TrimConfig,
90    pub inference: InferenceConfig,
91}
92
93impl MultiscreenConfig {
94    /// Validates all nested configuration.
95    pub fn validate(&self) -> Result<()> {
96        self.screens.validate()?;
97        self.tiles.validate()?;
98        self.trim.validate()?;
99        self.inference.validate()?;
100        Ok(())
101    }
102
103    /// Returns a compact config useful for examples and fast tests.
104    pub fn tiny() -> Self {
105        Self {
106            screens: ScreenConfig {
107                tokens_per_screen: 8,
108                screen_stride_tokens: 4,
109                max_screen_count: None,
110            },
111            tiles: TileConfig {
112                tokens_per_tile: 4,
113                tile_stride_tokens: 2,
114                screening_grid: crate::tile::ScreeningGridConfig {
115                    layer_count: 1,
116                    head_count: 2,
117                },
118            },
119            trim: TrimConfig {
120                acceptance_sharpness_r: 2.0,
121                screening_window_w: 4.0,
122            },
123            inference: InferenceConfig::default(),
124        }
125    }
126}