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
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
80
81
82
83
84
85
86
87
88
89
90
91
92

use syntect::easy::HighlightLines;
use syntect::highlighting::{Style, ThemeSet};
use syntect::parsing::{SyntaxReference, SyntaxSet};
use syntect::util::{as_24_bit_terminal_escaped, LinesWithEndings};

pub struct CodeHighlighter<'a> {
    syntax_set: &'a SyntaxSet,
    highlighter: HighlightLines<'a>,
}

impl<'a> CodeHighlighter<'a> {
    pub fn new(
        syntax_set: &'a SyntaxSet,
        syntax: &'a SyntaxReference,
        theme: &'a syntect::highlighting::Theme,
    ) -> Self {
        Self {
            syntax_set,
            highlighter: HighlightLines::new(syntax, theme),
        }
    }

    pub fn write(&mut self, code: &str) {
        for line in LinesWithEndings::from(code) {
            let ranges: Vec<(Style, &str)> = self
                .highlighter
                .highlight_line(line, self.syntax_set)
                .unwrap_or_default();
            let escaped = as_24_bit_terminal_escaped(&ranges[..], true);
            print!("{}", escaped);
        }
        print!("\x1b[0m");
    }

    pub fn finish_line(&mut self) {
        println!();
    }
}

pub fn resolve_syntax<'a>(syntax_set: &'a SyntaxSet, language: &str) -> &'a SyntaxReference {
    let lang = language.trim();
    if lang.is_empty() {
        return syntax_set.find_syntax_plain_text();
    }

    let lang_lc = lang.to_ascii_lowercase();
    let ext = match lang_lc.as_str() {
        "py" | "python" => "py",
        "rs" | "rust" => "rs",
        "js" | "javascript" => "js",
        "ts" | "typescript" => "ts",
        "sh" | "bash" | "shell" => "sh",
        "yml" | "yaml" => "yml",
        "md" | "markdown" => "md",
        "json" => "json",
        "toml" => "toml",
        _ => lang_lc.as_str(),
    };

    syntax_set
        .find_syntax_by_extension(ext)
        .or_else(|| syntax_set.find_syntax_by_extension(&lang_lc))
        .or_else(|| {
            syntax_set
                .syntaxes()
                .iter()
                .find(|s| s.name.to_ascii_lowercase() == lang_lc)
        })
        .unwrap_or_else(|| syntax_set.find_syntax_plain_text())
}

pub fn resolve_theme(theme_set: &ThemeSet) -> &syntect::highlighting::Theme {
    let theme_name = if theme_set.themes.contains_key("Dark+ (default dark)") {
        "Dark+ (default dark)"
    } else if theme_set.themes.contains_key("Monokai Extended") {
        "Monokai Extended"
    } else if theme_set.themes.contains_key("base16-ocean.dark") {
        "base16-ocean.dark"
    } else {
        theme_set
            .themes
            .keys()
            .next()
            .map(|s| s.as_str())
            .unwrap_or("base16-ocean.dark")
    };

    theme_set
        .themes
        .get(theme_name)
        .unwrap_or_else(|| &theme_set.themes["base16-ocean.dark"])
}