use crate::prelude::*;
#[derive(Debug)]
pub struct Template
{
#[allow(dead_code)]
source: String,
}
impl Template
{
pub fn new(source: impl Into<String>) -> Result<Template>
{
let source = source.into();
Ok(Template { source })
}
#[allow(dead_code)]
pub(crate) fn render(&self, messages: &[crate::Message]) -> Result<String>
{
let mut env = minijinja::Environment::new();
env.add_template("template", &self.source)?;
Ok(
env
.get_template("template")?
.render(minijinja::context! { messages => messages })?,
)
}
}
#[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()
}]).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"
);
}
}