Part II: Voice Assistant Desktop App with LLaMA3 and Whisper in Rust

 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110

/// This tauri command would receive a Vec<f32> which represents a chunk of audio being recorded
/// The chunk will be forwarded through the MPSC channel
#[tauri::command]
pub fn audio_chunk(
    app: tauri::State<'_, Arc<Instruct>>,
    req: ipc::Request<'_>
) -> Result<(), &'static str> {
    if let tauri::ipc::InvokeBody::Raw(data) = req.body() {
        let chunk = bytes_to_f32(&data[..]);

        if let Err(e) = app.send(chunk) {
            error!("audio_chunk: error: {e:?}");
            return Err("invalid chunk")    
        }
    } else {
        return Err("invalid chunk")
    }
    Ok(())
}

37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55

impl Instruct {
    // .. code ommitted ..

    /// Exposes an API to send data into our MPSC channel
    pub fn send(&self, data: Vec<f32>) -> Result<()> {
        self.send.send(data)?;
        
        Ok(())
    }

    // This method just forwards the incoming chunks to a method exposed by our `struct WhisperWrap`. The client doesn't need to wait for this to happen
    fn listen(app: Arc<Instruct>, recv: Receiver<Vec<f32>>) {
        while let Ok(next) = recv.recv() {
            app.whisper.chunk(next);
        }
    }

    // .. code ommitted ..
}

265
266
267
268
269
270
271
272
273
274
275
276
277

impl WhisperWrap {
    // ... code ommitted ...

    /// Accepts an incoming chunk of data and appends it to our `data` field of the struct
    pub fn chunk(&self, chunk: Vec<f32>) {
        let mut c = chunk;
        let mut chunk = self.data.lock().unwrap();
        // while appending this also `drains` the incoming Vec<>, saving some space
        chunk.append(&mut c);
    }

    // ... code ommitted ...
}

59
60
61
62
63
64
65
66
67
68

/// Public API to trigger audio inference
pub fn audio(&self) -> Result<Response> {
    let (transcript, n_tokens, elapsed) = self.whisper.infer()?;

    // More when we work on the `Pipeline` part of our inference

    Ok(
        Response::new(&transcript, &transcript, n_tokens, 0)
    )
}

143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257

impl WhisperWrap {
    // ... code ommitted ..
    
    // 1. Checks if we have valid data
    // 2. Creates the `Mel Spectogram` representation of our audio data
    // 3. Creates and returns a `Tensor` from the given data
    fn preproc(&self) -> Result<Tensor> {
        let data = match self.data.lock() {
            Ok(mut d) => {
                if d.len() < 4096 * 4 {
                    anyhow::bail!("Not enough audio data in buffer!");
                }
                let d = d.drain(..).collect::<Vec<_>>();
                
                d
            }
            Err(e) => {
                error!("error acquiring data lock: {e:?}");
                anyhow::bail!("Not enough audio data in buffer!");
            }
        };

        let mel = pcm_to_mel(&self.config, &data[..], &self.mel_filters[..]);
        let mel_len = mel.len();
        let mel = Tensor::from_vec(
            mel,
            (1, self.config.num_mel_bins, mel_len / self.config.num_mel_bins),
            &self.device,
        )?;

        Ok(mel)
    }

    // ... code ommitted ..

    /// Runs transcription
    pub fn infer(&self) -> Result<(String, u32, std::time::Duration)> {
        // generates `mel`
        let mels = self.preproc()?;

        let mut model = match self.model.lock() {
            Ok(m) => m,
            Err(e) => {
                error!("infer: error acquiring model lock: {e:?}");
                anyhow::bail!("error during inference");
            }
        };

        let (_, _, content_frames) = mels.dims3()?;
        let mut seek = 0;
        let mut segments = vec![];

        // newline tokens after each segment
        let nltokens = self.tokenizer.encode("\n", false).unwrap().get_ids().to_vec();

        let mut total_dur = Duration::from_millis(0);
        let mut total_tokens = 0;

        // seek through the generated `mels` and call the `decode_segment` method on a chunk
        while seek < content_frames {
            let start = std::time::Instant::now();

            let segment_size = usize::min(content_frames - seek, N_FRAMES);
            let mel_segment = mels.narrow(2, seek, segment_size)?;

            let mut decoded = self.decode_segment(&mut model, &mel_segment)?;
            seek += segment_size;

            total_dur += std::time::Instant::now() - start;
            total_tokens += decoded.tokens.len();

            if decoded.no_speech_prob > NO_SPEECH_THRESHOLD && decoded.avg_logprob < LOGPROB_THRESHOLD {
                println!("no speech detected, skipping {seek} {decoded:?}");
                continue;
            }
            
            segments.append(&mut decoded.tokens);
            // adding newline tokens after each segment
            nltokens.iter().for_each(|&t| {
                segments.push(t);
            });
        }

        // Let us now create the final text output
        let instruct = self.tokenizer.decode(&segments, true).map_err(|_| anyhow!("error creating text from tokens"))?;

        Ok((instruct, total_tokens as u32, total_dur))
    }

    // Decodes a single segment at different `hyperparameter temperature`
    // The current values of the hyperparameter TEMPERATURES = [0.0, 0.2, 0.4, 0.6, 0.8, 1.0]
    fn decode_segment(&self, model: &mut Whisper, segment: &Tensor) -> Result<DecodingResult> {
        // Decode at a particular temperature, check if we have a valid result or move on to the next temperature
        for (i, &t) in TEMPERATURES.iter().enumerate() {
            let decoded = self.decode(model, segment, t);
            if i == TEMPERATURES.len() - 1 {
                return decoded;
            }

            match decoded {
                Ok(decoded) => {
                    if decoded.avg_logprob >= LOGPROB_THRESHOLD || decoded.no_speech_prob > NO_SPEECH_THRESHOLD {
                        info!("Decoded: {decoded:?}");
                        return Ok(decoded);
                    } 
                }
                Err(e) => {
                    warn!("Error decoding @ temperature: {t}: {e:?}");
                }
            }
        }

        unreachable!()
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79

// inside impl WhisperWrap
fn decode(&self, model: &mut Whisper, segment: &Tensor, temp: f64) -> Result<DecodingResult> {
    // generating some random seed
    let mut rng = rand::thread_rng();

    // the `mel` segment block is passed through the encoder here, you can think of this process as generating `features` from your input audio data segment
    let features = model.encoder.forward(segment, true)?;

    // some token pre-generation, basically creating a `token` representation which tells the `decode` part of the model [start of transcript, english, task transcribe, don't generate timestamps] - these special tokens kind of act like a prompt for the decoder model
    let mut tokens = self.preproc_decode();

    // now, we initialize some variables that will maintain our stats and metrics derived from this part of the inference
    // probability of this segment being some random/ background noise
    let mut no_speech_prob = f64::MAX;
    // the average probability of our prediction across all the decoding passes
    let mut sum_log_p = 0.;

    // Ok, so now we loop through a `max number of possible outputs` and `autoregressively` generate the next token
    for i in 0..self.config.max_target_positions {
        // we'll comvert the `Vec<tokens>` slice to a `Tensor`
        let tensor = Tensor::new(&tokens[..], &self.device)?;
        // and send that tensor to the `decoder` to generate the next token
        let dec = model.decoder.forward(&tensor.unsqueeze(0)?, &features, i == 0)?;

        // Extract the no speech probability on the first iteration by looking at the first
        // token logits and the probability for the according token.
        if i == 0 {
            let logits = model.decoder.final_linear(&dec.i(..1)?)?.i(0)?.i(0)?;
            no_speech_prob = softmax(&logits, 0)?
                .i(self.default_tokens.no_speech as usize)?
                .to_scalar::<f32>()? as f64;
        }

        let (_, seq_len, _) = dec.dims3()?;
        let logits = model.decoder.final_linear(&dec.i((..1, seq_len - 1..))?)?
            .i(0)?
            .i(0)?;

        // a simple sampler, picked up from `https://github.com/huggingface/candle/blob/main/candle-examples/examples/whisper/main.rs`
        let next_token = if temp > 0. {
            let prs = softmax(&(&logits / temp)?, 0)?;
            let logits_v: Vec<f32> = prs.to_vec1()?;
            let distr = rand::distributions::WeightedIndex::new(&logits_v)?;
            distr.sample(&mut rng) as u32
        } else {
            let logits_v: Vec<f32> = logits.to_vec1()?;
            logits_v
                .iter()
                .enumerate()
                .max_by(|(_, u), (_, v)| u.total_cmp(v))
                .map(|(i, _)| i as u32)
                .unwrap()
        };

        // remember, that this decoding is `autoregressive`, meaning, output of the current pass is passed on as input to the next pass till some stop condition is reached
        tokens.push(next_token);

        // bookkeeping the `probability`, we'll calculate average on this later and that would serve as our `decisioning` benchmark comparing with `LOGPROB_THRESHOLD`
        let prob = softmax(&logits, candle_core::D::Minus1)?
            .i(next_token as usize)?
            .to_scalar::<f32>()? as f64;

        // the stop condition, if end-of-token or max target has been reached
        if next_token == self.default_tokens.eot || tokens.len() > self.config.max_target_positions {
            break;
        }
        sum_log_p += prob.ln();
    }
    
    // Finally create a struct to hold our output and metadata for this decoding pass
    Ok(DecodingResult {
        text: self.tokenizer.decode(&tokens, true).map_err(|_| anyhow!("error creating text from tokens"))?,
        avg_logprob: sum_log_p / tokens.len() as f64,
        tokens,
        no_speech_prob,
        temperature: temp,

    })
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22

<script lang="ts">
    // Adapted from https://github.com/kgullion/vite-typescript-audio-worklet-example/blob/main/src/main.ts
    import audioProcUrl from "$lib/audio-proc/audio-processor?url";
    import { invoke } from '@tauri-apps/api/core';
    import type { Inference, QuestionAnswer } from "$lib/types";
    import Qa from "./QA.svelte";

    // defining a constant buffer size for chunking audio
    const BUFFER_SIZE = 4096;
    // sampling rate for whisper 16KHz
    const SAMPLE_RATE = 16000; // Whisper typically expects 16kHz audio

    // ... code ommitted ...

    // a variable to hold the media stream
    let stream: MediaStream|null = null;

    // an audio context 
    let audioContext: AudioContext|null = null;
    let source: MediaStreamAudioSourceNode|null = null;
    let workletNode: AudioWorkletNode|null = null;
</script>

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36

// Adapted from https://github.com/kgullion/vite-typescript-audio-worklet-example/blob/main/src/main.ts

class AudioProcessor extends AudioWorkletProcessor {
    private bufferSize: number;
    private buffer: Float32Array;
    private bufferIndex: number;

    constructor(options?: AudioWorkletNodeOptions) {
        super();
        this.bufferSize = options?.processorOptions.bufferSize || 4096;
        this.buffer = new Float32Array(this.bufferSize);
        this.bufferIndex = 0;
    }

    // this method is basically responsible for an input buffer and creating a output chunked buffer of specific size
    process(inputs: Float32Array[][], outputs: Float32Array[][], parameters: Record<string, Float32Array>): boolean {
        const input = inputs[0];
        const channel = input[0];

        if (channel) {
            for (let i = 0; i < channel.length; i++) {
                this.buffer[this.bufferIndex++] = channel[i];

                if (this.bufferIndex === this.bufferSize) {
                    this.port.postMessage(this.buffer);
                    this.buffer = new Float32Array(this.bufferSize);
                    this.bufferIndex = 0;
                }
            }
        }

        return true;
    }
}

registerProcessor('audio-processor', AudioProcessor);

64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80

impl Instruct {
    // .. code ommitted ..


    /// Public API to trigger audio inference
    pub fn audio(&self) -> Result<Response> {
        let (transcript, n_tokens, elapsed) = self.whisper.infer()?;
        let (generated, n_txt_tok, txt_elapsed) = self.llama.infer(&transcript)?;

        Ok(Response::new(
            &transcript,
            &generated,
            (n_tokens + n_txt_tok) as u32,
            (elapsed + txt_elapsed).as_secs(),
        ))
    }
}