ravenrs_ai/
engine.rs

1//! Runs a model on this machine, on the CPU.
2
3use std::fmt;
4use std::fs::File;
5use std::time::Instant;
6
7use candle_core::quantized::gguf_file;
8use candle_core::{Device, Tensor};
9use candle_transformers::generation::{LogitsProcessor, Sampling};
10use candle_transformers::models::{
11    quantized_gemma3, quantized_llama, quantized_phi3, quantized_qwen2, quantized_qwen3,
12};
13use tokenizers::Tokenizer;
14
15use crate::{LocalModel, ModelFamily, ModelFolder, lfm3};
16
17/// How sure it answers: low, because it should stick to what it is given.
18const TEMPERATURE: f64 = 0.2;
19const SEED: u64 = 42;
20
21/// The weights of each kind of model.
22enum Weights {
23    Qwen3(quantized_qwen3::ModelWeights),
24    Qwen2(quantized_qwen2::ModelWeights),
25    Gemma3(quantized_gemma3::ModelWeights),
26    Llama(quantized_llama::ModelWeights),
27    Phi3(quantized_phi3::ModelWeights),
28    Lfm3(lfm3::ModelWeights),
29}
30
31impl Weights {
32    /// The logits for the next token. At offset 0 each kind forgets what it read before.
33    fn forward(&mut self, input: &Tensor, offset: usize) -> candle_core::Result<Tensor> {
34        let logits = match self {
35            Self::Qwen3(model) => model.forward(input, offset),
36            Self::Qwen2(model) => model.forward(input, offset),
37            Self::Gemma3(model) => model.forward(input, offset),
38            Self::Llama(model) => model.forward(input, offset),
39            Self::Phi3(model) => model.forward(input, offset),
40            Self::Lfm3(model) => model.forward(input, offset),
41        }?;
42        // Some give one row per token read, others only the last.
43        let logits = logits.squeeze(0)?;
44        if logits.rank() == 2 {
45            let last = logits.dim(0)? - 1;
46            logits.get(last)
47        } else {
48            Ok(logits)
49        }
50    }
51
52    fn clear(&mut self) {
53        match self {
54            Self::Qwen3(model) => model.clear_kv_cache(),
55            Self::Qwen2(model) => model.clear_kv_cache(),
56            Self::Llama(model) => model.clear_kv_cache(),
57            Self::Lfm3(model) => model.clear_kv_cache(),
58            Self::Gemma3(_) | Self::Phi3(_) => {}
59        }
60    }
61}
62
63/// A family's way of writing a conversation, and the tokens that end its answer.
64pub(crate) fn prompt_of(family: ModelFamily, system: &str, user: &str) -> (String, &'static [&'static str]) {
65    match family {
66        ModelFamily::Qwen3 => (
67            format!(
68                "<|im_start|>system\n{system}<|im_end|>\n<|im_start|>user\n{user} /no_think<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n"
69            ),
70            &["<|im_end|>", "<|endoftext|>"],
71        ),
72        ModelFamily::Qwen2 => (
73            format!(
74                "<|im_start|>system\n{system}<|im_end|>\n<|im_start|>user\n{user}<|im_end|>\n<|im_start|>assistant\n"
75            ),
76            &["<|im_end|>", "<|endoftext|>"],
77        ),
78        ModelFamily::Gemma3 => (
79            format!("<bos><start_of_turn>user\n{system}\n\n{user}<end_of_turn>\n<start_of_turn>model\n"),
80            &["<end_of_turn>", "<eos>"],
81        ),
82        ModelFamily::Llama3 => (
83            format!(
84                "<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\n{system}<|eot_id|><|start_header_id|>user<|end_header_id|>\n\n{user}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
85            ),
86            &["<|eot_id|>", "<|end_of_text|>"],
87        ),
88        ModelFamily::Phi3 => (
89            format!("<|system|>\n{system}<|end|>\n<|user|>\n{user}<|end|>\n<|assistant|>\n"),
90            &["<|end|>", "<|endoftext|>"],
91        ),
92        ModelFamily::Mistral => (format!("<s>[INST] {system}\n\n{user} [/INST]"), &["</s>"]),
93        ModelFamily::Lfm3 => (
94            format!(
95                "<|startoftext|><|im_start|>system\n{system}<|im_end|>\n<|im_start|>user\n{user}<|im_end|>\n<|im_start|>assistant\n<think></think>"
96            ),
97            &["<|im_end|>"],
98        ),
99        // It does not converse: a `Transcriber` runs it.
100        ModelFamily::Whisper | ModelFamily::Kokoro => (String::new(), &[]),
101    }
102}
103
104/// Why a model did not run.
105#[derive(Clone, Debug, PartialEq, Eq)]
106#[non_exhaustive]
107pub enum EngineError {
108    /// Its files are missing or are not what this engine reads.
109    Load(String),
110    /// It stopped in the middle.
111    Run(String),
112}
113
114impl fmt::Display for EngineError {
115    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
116        match self {
117            Self::Load(why) => write!(f, "the model could not be loaded: {why}"),
118            Self::Run(why) => write!(f, "the model stopped: {why}"),
119        }
120    }
121}
122
123impl std::error::Error for EngineError {}
124
125#[cfg(test)]
126mod tests {
127    use super::prompt_of;
128    use crate::ModelFamily;
129
130    #[test]
131    fn each_family_writes_the_conversation_its_way() {
132        let families = [
133            ModelFamily::Qwen3,
134            ModelFamily::Qwen2,
135            ModelFamily::Gemma3,
136            ModelFamily::Llama3,
137            ModelFamily::Phi3,
138            ModelFamily::Mistral,
139            ModelFamily::Lfm3,
140        ];
141        for family in families {
142            let (prompt, ends) = prompt_of(family, "RULES", "QUESTION");
143            assert!(prompt.find("RULES") < prompt.find("QUESTION"), "{family:?}");
144            assert!(!ends.is_empty(), "{family:?}");
145        }
146        assert!(prompt_of(ModelFamily::Qwen3, "", "").0.contains("/no_think"));
147        assert!(!prompt_of(ModelFamily::Qwen2, "", "").0.contains("/no_think"));
148    }
149}
150
151/// What a generation cost.
152#[derive(Clone, Copy, Debug, PartialEq)]
153pub struct Generated {
154    /// Tokens it wrote.
155    pub tokens: usize,
156    /// Seconds until the first of them.
157    pub first_token: f64,
158    /// Seconds in all.
159    pub seconds: f64,
160}
161
162impl Generated {
163    /// Tokens per second after the first one.
164    pub fn tokens_per_second(&self) -> f64 {
165        let writing = self.seconds - self.first_token;
166        if self.tokens > 1 && writing > 0.0 { (self.tokens - 1) as f64 / writing } else { 0.0 }
167    }
168}
169
170/// A model loaded in memory. Loading takes a while and generating blocks, so keep it
171/// away from the UI thread.
172pub struct LocalEngine {
173    family: ModelFamily,
174    model: Weights,
175    tokenizer: Tokenizer,
176    device: Device,
177    ends: Vec<u32>,
178}
179
180impl fmt::Debug for LocalEngine {
181    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
182        f.debug_struct("LocalEngine").field("device", &self.device).finish_non_exhaustive()
183    }
184}
185
186impl LocalEngine {
187    /// Loads `model` from `folder`.
188    ///
189    /// # Errors
190    ///
191    /// When its files are missing or cannot be read.
192    pub fn load(folder: &ModelFolder, model: &LocalModel) -> Result<Self, EngineError> {
193        let load = |why: String| EngineError::Load(why);
194        let named = |name: &str| {
195            let file = model.files.iter().find(|file| file.name == name);
196            file.map(|file| folder.path_of(model, file)).ok_or_else(|| load(format!("no {name}")))
197        };
198        let device = Device::Cpu;
199        let mut file = File::open(named("model.gguf")?).map_err(|e| load(e.to_string()))?;
200        let content = gguf_file::Content::read(&mut file).map_err(|e| load(e.to_string()))?;
201        let family = model.family;
202        let weights = match family {
203            ModelFamily::Qwen3 => {
204                quantized_qwen3::ModelWeights::from_gguf(content, &mut file, &device).map(Weights::Qwen3)
205            }
206            ModelFamily::Qwen2 => {
207                quantized_qwen2::ModelWeights::from_gguf(content, &mut file, &device).map(Weights::Qwen2)
208            }
209            ModelFamily::Gemma3 => {
210                quantized_gemma3::ModelWeights::from_gguf(content, &mut file, &device).map(Weights::Gemma3)
211            }
212            ModelFamily::Llama3 | ModelFamily::Mistral => {
213                quantized_llama::ModelWeights::from_gguf(content, &mut file, &device).map(Weights::Llama)
214            }
215            ModelFamily::Phi3 => {
216                quantized_phi3::ModelWeights::from_gguf(false, content, &mut file, &device).map(Weights::Phi3)
217            }
218            ModelFamily::Lfm3 => {
219                lfm3::ModelWeights::from_gguf(&content, &mut file, &device).map(Weights::Lfm3)
220            }
221            ModelFamily::Whisper => return Err(load("Whisper hears; a Transcriber runs it".into())),
222            ModelFamily::Kokoro => return Err(load("Kokoro speaks; a SpeechEngine runs it".into())),
223        }
224        .map_err(|e| load(e.to_string()))?;
225        let tokenizer = Tokenizer::from_file(named("tokenizer.json")?).map_err(|e| load(e.to_string()))?;
226        let ends = prompt_of(family, "", "").1.iter().filter_map(|end| tokenizer.token_to_id(end)).collect();
227        Ok(Self { family, model: weights, tokenizer, device, ends })
228    }
229
230    /// Answers `user`, following `system`, in at most `max_tokens`. `on_text` gets the
231    /// text as it comes, and returning false stops it.
232    ///
233    /// # Errors
234    ///
235    /// When it stops in the middle.
236    pub fn generate(
237        &mut self,
238        system: &str,
239        user: &str,
240        max_tokens: usize,
241        on_text: &mut dyn FnMut(&str) -> bool,
242    ) -> Result<Generated, EngineError> {
243        let run = |why: String| EngineError::Run(why);
244        let (prompt, _) = prompt_of(self.family, system, user);
245        let started = Instant::now();
246        self.model.clear();
247        let encoding = self.tokenizer.encode(prompt, false).map_err(|e| run(e.to_string()))?;
248        let prompt_tokens = encoding.get_ids();
249        let mut sampler = LogitsProcessor::from_sampling(SEED, Sampling::All { temperature: TEMPERATURE });
250        let mut written: Vec<u32> = Vec::new();
251        let mut shown = 0;
252        let mut first_token = 0.0;
253        let mut input = prompt_tokens.to_vec();
254        let mut offset = 0;
255        while written.len() < max_tokens {
256            let tensor = Tensor::new(input.as_slice(), &self.device)
257                .and_then(|t| t.unsqueeze(0))
258                .map_err(|e| run(e.to_string()))?;
259            let logits = self.model.forward(&tensor, offset).map_err(|e| run(e.to_string()))?;
260            let next = sampler.sample(&logits).map_err(|e| run(e.to_string()))?;
261            if written.is_empty() {
262                first_token = started.elapsed().as_secs_f64();
263            }
264            offset += input.len();
265            input = vec![next];
266            if self.ends.contains(&next) {
267                break;
268            }
269            written.push(next);
270            let text = self.tokenizer.decode(&written, true).map_err(|e| run(e.to_string()))?;
271            // A letter split between tokens shows up only when it is whole.
272            if !text.ends_with('\u{fffd}') && text.len() > shown {
273                let keep_going = on_text(text.get(shown..).unwrap_or_default());
274                shown = text.len();
275                if !keep_going {
276                    break;
277                }
278            }
279        }
280        Ok(Generated { tokens: written.len(), first_token, seconds: started.elapsed().as_secs_f64() })
281    }
282}
ravenrs_ai/engine.rs

ravenrs_ai/
engine.rs
