kproc-llm 0.7.0

Knowledge Processing library, using LLMs.
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

//! Template module
use std::collections::HashMap;

use crate::prelude::*;

/// Template
#[derive(Debug)]
pub struct Template
{
  #[allow(dead_code)]
  source: String,
  add_generation_prompt: bool,
}

impl Template
{
  /// New template
  pub fn new(source: impl Into<String>) -> Result<Template>
  {
    let source = source.into();
    Ok(Template {
      source,
      add_generation_prompt: true,
    })
  }
  #[allow(dead_code)]
  pub(crate) fn render(
    &self,
    messages: &[crate::Message],
    enable_thinking: bool,
    options: HashMap<String, minijinja::Value>,
  ) -> Result<String>
  {
    let mut env = minijinja::Environment::new();
    env.add_template("template", &self.source)?;
    let rendered = env.get_template("template")?.render(minijinja::context! {
      messages => messages,
      add_generation_prompt => self.add_generation_prompt,
      enable_thinking => enable_thinking,
      ..options
    })?;
    Ok(rendered)
  }
}

#[cfg(test)]
mod tests
{
  #[test]
  fn template_llama()
  {
    let template = super::Template::new(include_str!("../data/templates/llama")).unwrap();
    assert_eq!(
      template.render(&[crate::Message {
        role: crate::Role::User,
        content: "Hello world.".into()
      }], false, Default::default()).unwrap(),
      "<|start_header_id|>user<|end_header_id|>Hello world.<|eot_id|>\n<|start_header_id|>assistant<|end_header_id|><|start_header_id|>assistant<|end_header_id|>\n"
    );
  }
}