slm_inference 0.1.1

Backend-agnostic trait layer for running Small Language Model (SLM) inference in Rust.
Documentation 
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

use crate::SlmRole;
use crate::formatter::{SlmFormatter, SlmToolStyle};

pub struct Llama3Formatter;

impl SlmFormatter for Llama3Formatter {
    fn bos(&self) -> Option<&str> {
        Some("<|begin_of_text|>")
    }
    fn tool_style(&self) -> SlmToolStyle {
        SlmToolStyle::SeparateTurn
    }

    fn turn_start(&self, role: &SlmRole) -> String {
        match role {
            SlmRole::System => "<|start_header_id|>system<|end_header_id|>\n\n".to_string(),
            SlmRole::User => "<|start_header_id|>user<|end_header_id|>\n\n".to_string(),
            SlmRole::Assistant => "<|start_header_id|>assistant<|end_header_id|>\n\n".to_string(),
            SlmRole::Tool(_) => "<|start_header_id|>ipython<|end_header_id|>\n\n".to_string(),
        }
    }

    fn turn_end(&self, _role: &SlmRole) -> String {
        "<|eot_id|>".to_string()
    }

    fn reasoning_bounds(&self) -> Option<(&str, &str)> {
        Some(("<think>\n", "\n</think>\n"))
    }

    fn wrap_reasoning(&self, content: &str) -> String {
        format!("<think>\n{}\n</think>\n", content.trim())
    }

    fn reasoning_trigger(&self) -> Option<&str> {
        Some("\n<think>\n")
    }

    fn format_tool_call(&self, name: &str, arguments_json: &str) -> String {
        format!(
            "{{\"name\": \"{}\", \"parameters\": {}}}",
            name,
            arguments_json.trim()
        )
    }

    fn format_tool_response(&self, _name: &str, content: &str) -> String {
        format!("{}", content.trim())
    }

    fn strip_tags(&self, text: &str) -> String {
        let mut cleaned = text.to_string();

        let full_headers = [
            "<|begin_of_text|>",
            "<|end_of_text|>",
            "<|start_header_id|>system<|end_header_id|>\n\n",
            "<|start_header_id|>user<|end_header_id|>\n\n",
            "<|start_header_id|>assistant<|end_header_id|>\n\n",
            "<|start_header_id|>ipython<|end_header_id|>\n\n",
            "<|eot_id|>",
            "<|eom_id|>",
        ];

        for tag in full_headers {
            cleaned = cleaned.replace(tag, "");
        }

        cleaned
    }
}