model-rs 0.1.0

A Rust CLI tool for downloading HuggingFace models and running local LLM inference
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

use crate::error::Result;
use tokenizers::Tokenizer;

pub fn get_eos_token(tokenizer: &Tokenizer) -> Option<u32> {
    tokenizer.token_to_id("</s>")
        .or_else(|| tokenizer.token_to_id("<EOS>"))
}

pub fn stream_piece(tokenizer: &Tokenizer, token_id: u32, started: &mut bool) -> Result<Option<String>> {
    if let Some(raw) = tokenizer.id_to_token(token_id) {
        if raw == "</s>" || raw == "<s>" || raw == "<unk>" {
            return Ok(None);
        }

        let text = tokenizer.decode(&[token_id], false)?;
        if text.is_empty() {
            return Ok(None);
        }

        let mut piece = String::new();
        if raw.starts_with('▁') && *started {
            piece.push(' ');
        }
        piece.push_str(&text);
        *started = true;
        return Ok(Some(piece));
    }

    Ok(None)
}