Include system prompt

gabrielmbmb · Sep 2, 2024 · 4ec8537 · 4ec8537
1 parent 99e075d
commit 4ec8537
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/candle-holder-examples/examples/streamed_text_generation/main.rs b/candle-holder-examples/examples/streamed_text_generation/main.rs
@@ -46,8 +46,12 @@ fn main() -> Result<()> {
     let model = AutoModelForCausalLM::from_pretrained(args.model, &device, None, None)?;
 
     let mut encodings = if args.apply_chat_template {
+        let mut messages = vec![Message::user(args.prompt.clone())];
+        if let Some(system_prompt) = args.system_prompt {
+            messages.insert(0, Message::system(system_prompt))
+        }
         tokenizer
-            .apply_chat_template_and_encode(vec![Message::user(args.prompt)], true)
+            .apply_chat_template_and_encode(messages, true)
             .map_err(Error::msg)?
     } else {
         tokenizer