Skip to main content

kernel/profiles/
profile.rs

1//! The parameter-schema system: which tunable parameters a model exposes,
2//! assembled from capability- and runtime-specific profiles.
3
4use crate::records::{Capability, JsonValue, ModelRecord, ParamSpec, ParamType, RuntimeId};
5
6type Matcher = Box<dyn Fn(&ModelRecord) -> bool + Send + Sync>;
7
8/// A set of parameter specs that applies to any model the `matches` predicate
9/// accepts.
10pub struct ModelProfile {
11    /// A stable identifier for the profile.
12    pub id: String,
13    /// The parameter specs this profile contributes.
14    pub schema: Vec<ParamSpec>,
15    matcher: Matcher,
16}
17
18impl ModelProfile {
19    /// Build a profile from an id, its schema, and a predicate.
20    pub fn new(
21        id: &str,
22        schema: Vec<ParamSpec>,
23        matcher: impl Fn(&ModelRecord) -> bool + Send + Sync + 'static,
24    ) -> Self {
25        Self {
26            id: id.to_owned(),
27            schema,
28            matcher: Box::new(matcher),
29        }
30    }
31
32    /// Whether this profile applies to `record`.
33    pub fn matches(&self, record: &ModelRecord) -> bool {
34        (self.matcher)(record)
35    }
36}
37
38/// A collection of profiles that together determine a model's parameter schema.
39pub struct ProfileRegistry {
40    /// The profiles considered, in priority order.
41    pub profiles: Vec<ModelProfile>,
42}
43
44impl ProfileRegistry {
45    /// Build a registry from a list of profiles.
46    pub fn new(profiles: Vec<ModelProfile>) -> Self {
47        Self { profiles }
48    }
49
50    /// The parameter schema for `record`: every matching profile's specs (first
51    /// occurrence of each key wins), plus a `context_length` spec when the model
52    /// and runtime honor one.
53    pub fn schema(&self, record: &ModelRecord) -> Vec<ParamSpec> {
54        let mut specs: Vec<ParamSpec> = Vec::new();
55        for profile in &self.profiles {
56            if profile.matches(record) {
57                for spec in &profile.schema {
58                    if !specs.iter().any(|kept| kept.key == spec.key) {
59                        specs.push(spec.clone());
60                    }
61                }
62            }
63        }
64        if !specs.iter().any(|spec| spec.key == "context_length")
65            && let Some(context) = context_length_spec(record)
66        {
67            specs.push(context);
68        }
69        specs
70    }
71
72    /// A copy of `record` with its parameter schema refreshed. An empty schema
73    /// leaves the record unchanged.
74    pub fn refreshed(&self, record: &ModelRecord) -> ModelRecord {
75        let schema = self.schema(record);
76        if schema.is_empty() {
77            return record.clone();
78        }
79        let mut updated = record.clone();
80        updated.params = schema;
81        updated
82    }
83
84    /// The built-in profile set covering text generation, per-runtime sampling
85    /// extras, speech, transcription, and togglable thinking.
86    pub fn builtin() -> Self {
87        Self::new(vec![
88            ModelProfile::new(
89                "text-generation",
90                vec![temperature(), top_p(), max_tokens()],
91                is_text,
92            ),
93            runtime_extras(
94                "sampling-llama-cpp",
95                RuntimeId::llama_cpp(),
96                vec![
97                    top_k(),
98                    min_p(),
99                    repeat_penalty(),
100                    frequency_penalty(),
101                    presence_penalty(),
102                    seed(),
103                    stop(),
104                ],
105            ),
106            runtime_extras(
107                "sampling-mlx-swift",
108                RuntimeId::mlx_swift(),
109                vec![repeat_penalty(), stop()],
110            ),
111            runtime_extras(
112                "sampling-mlx-lm",
113                RuntimeId::mlx_lm(),
114                vec![top_k(), min_p(), repeat_penalty(), seed(), stop()],
115            ),
116            runtime_extras(
117                "sampling-ollama",
118                RuntimeId::ollama(),
119                vec![
120                    top_k(),
121                    min_p(),
122                    seed(),
123                    repeat_penalty(),
124                    frequency_penalty(),
125                    presence_penalty(),
126                    stop(),
127                ],
128            ),
129            runtime_extras(
130                "sampling-endpoint",
131                RuntimeId::openai_endpoint(),
132                vec![stop(), seed(), frequency_penalty(), presence_penalty()],
133            ),
134            runtime_extras(
135                "sampling-apple",
136                RuntimeId::apple_foundation(),
137                vec![top_k(), seed()],
138            ),
139            ModelProfile::new(
140                "speech-synthesis",
141                vec![plain("voice", ParamType::String), speed()],
142                |record| record.can(&Capability::speak()),
143            ),
144            ModelProfile::new(
145                "transcription",
146                vec![
147                    plain("language", ParamType::String),
148                    plain("translate", ParamType::Bool),
149                ],
150                |record| record.can(&Capability::transcribe()),
151            ),
152            ModelProfile::new(
153                "togglable-thinking",
154                vec![plain("thinking", ParamType::Bool)],
155                |record| {
156                    record.can(&Capability::chat())
157                        && record.runtime.id.as_ref().is_some_and(is_thinking_runtime)
158                },
159            ),
160        ])
161    }
162}
163
164/// The `context_length` spec for a chat/completion model on a runtime that honors
165/// it (llama.cpp or Ollama), sized to the model's declared window.
166pub fn context_length_spec(record: &ModelRecord) -> Option<ParamSpec> {
167    if !(record.can(&Capability::chat()) || record.can(&Capability::complete())) {
168        return None;
169    }
170    let runtime = record.runtime.id.as_ref()?;
171    if !is_context_honoring(runtime) {
172        return None;
173    }
174    match record.context_length {
175        Some(window) if window > 0 => Some(ParamSpec {
176            key: "context_length".to_owned(),
177            param_type: ParamType::Int,
178            default_value: Some(JsonValue::Int(window.min(32768))),
179            range: Some(vec![
180                JsonValue::Int(512.min(window)),
181                JsonValue::Int(window),
182            ]),
183            values: None,
184        }),
185        _ => Some(ParamSpec {
186            key: "context_length".to_owned(),
187            param_type: ParamType::Int,
188            default_value: None,
189            range: Some(vec![JsonValue::Int(512), JsonValue::Int(131072)]),
190            values: None,
191        }),
192    }
193}
194
195fn is_text(record: &ModelRecord) -> bool {
196    record.can(&Capability::chat()) || record.can(&Capability::complete())
197}
198
199fn is_context_honoring(runtime: &RuntimeId) -> bool {
200    *runtime == RuntimeId::ollama() || *runtime == RuntimeId::llama_cpp()
201}
202
203fn is_thinking_runtime(runtime: &RuntimeId) -> bool {
204    *runtime == RuntimeId::ollama() || *runtime == RuntimeId::mlx_lm()
205}
206
207fn runtime_extras(id: &str, runtime: RuntimeId, schema: Vec<ParamSpec>) -> ModelProfile {
208    ModelProfile::new(id, schema, move |record| {
209        is_text(record) && record.runtime.id.as_ref() == Some(&runtime)
210    })
211}
212
213fn float(key: &str, low: f64, high: f64) -> ParamSpec {
214    ParamSpec {
215        key: key.to_owned(),
216        param_type: ParamType::Float,
217        default_value: None,
218        range: Some(vec![JsonValue::Double(low), JsonValue::Double(high)]),
219        values: None,
220    }
221}
222
223fn int(key: &str, low: i64, high: i64) -> ParamSpec {
224    ParamSpec {
225        key: key.to_owned(),
226        param_type: ParamType::Int,
227        default_value: None,
228        range: Some(vec![JsonValue::Int(low), JsonValue::Int(high)]),
229        values: None,
230    }
231}
232
233fn plain(key: &str, param_type: ParamType) -> ParamSpec {
234    ParamSpec {
235        key: key.to_owned(),
236        param_type,
237        default_value: None,
238        range: None,
239        values: None,
240    }
241}
242
243fn temperature() -> ParamSpec {
244    float("temperature", 0.0, 2.0)
245}
246fn top_p() -> ParamSpec {
247    float("top_p", 0.0, 1.0)
248}
249fn top_k() -> ParamSpec {
250    int("top_k", 0, 100)
251}
252fn min_p() -> ParamSpec {
253    float("min_p", 0.0, 1.0)
254}
255fn max_tokens() -> ParamSpec {
256    int("max_tokens", 1, 32768)
257}
258fn repeat_penalty() -> ParamSpec {
259    float("repeat_penalty", 0.5, 2.0)
260}
261fn frequency_penalty() -> ParamSpec {
262    float("frequency_penalty", -2.0, 2.0)
263}
264fn presence_penalty() -> ParamSpec {
265    float("presence_penalty", -2.0, 2.0)
266}
267fn seed() -> ParamSpec {
268    plain("seed", ParamType::Int)
269}
270fn stop() -> ParamSpec {
271    plain("stop", ParamType::String)
272}
273fn speed() -> ParamSpec {
274    ParamSpec {
275        key: "speed".to_owned(),
276        param_type: ParamType::Float,
277        default_value: Some(JsonValue::Double(1.0)),
278        range: Some(vec![JsonValue::Double(0.5), JsonValue::Double(2.0)]),
279        values: None,
280    }
281}