Pr 2395 ci run (#2406)

* fix(router): Fix appending to message content

* feat: add message and chat template test

---------

Co-authored-by: Simone Rossi <simone.rossi.93@gmail.com>
This commit is contained in:
drbh 2024-08-12 14:38:59 -04:00 committed by GitHub
parent 19ea85f8dc
commit 9a7830bd28
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 72 additions and 3 deletions

View File

@ -94,7 +94,9 @@ impl ChatTemplate {
mod tests { mod tests {
use crate::infer::chat_template::raise_exception; use crate::infer::chat_template::raise_exception;
use crate::infer::ChatTemplate; use crate::infer::ChatTemplate;
use crate::{ChatTemplateInputs, Message, MessageContent, TextMessage, TokenizerConfigToken}; use crate::{
ChatTemplateInputs, GrammarType, Message, MessageContent, TextMessage, TokenizerConfigToken,
};
use minijinja::Environment; use minijinja::Environment;
#[test] #[test]
@ -823,4 +825,40 @@ mod tests {
} }
} }
} }
#[test]
fn test_chat_template_with_default_tool_template() {
let ct = ChatTemplate::new(
"{{ bos_token }}{% for message in messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if message['role'] == 'user' %}{{ '[INST] ' + message['content'] + ' [/INST]' }}{% elif message['role'] == 'assistant' %}{{ message['content'] + eos_token + ' ' }}{% else %}{{ raise_exception('Only user and assistant roles are supported!') }}{% endif %}{% endfor %}".to_string(),
Some(TokenizerConfigToken::String("<s>".to_string())),
Some(TokenizerConfigToken::String("</s>".to_string())),
);
// convert TextMessage to Message
let msgs: Vec<Message> = vec![
Message {
name: None,
role: "user".to_string(),
content: MessageContent::SingleText(
"I'd like to show off how chat templating works!".to_string(),
),
},
Message {
name: None,
role: "assistant".to_string(),
content: MessageContent::SingleText("Great! How can I help you today?".to_string()),
},
Message {
name: None,
role: "user".to_string(),
content: MessageContent::SingleText("Just testing".to_string()),
},
];
let tools = serde_json::json!("[]");
let tool_prompt = "This default prompt will be used".to_string();
let grammer_with_prompt = (GrammarType::Json(tools), tool_prompt);
let result = ct.apply(None, msgs, Some(grammer_with_prompt));
let expected = "<s>[INST] I'd like to show off how chat templating works! [/INST]Great! How can I help you today?</s> [INST] Just testing\n---\nThis default prompt will be used\n\"[]\" [/INST]".to_string();
assert_eq!(result.unwrap(), expected);
}
} }

View File

@ -1016,8 +1016,10 @@ impl MessageContent {
pub fn push(&mut self, chunk: MessageChunk) { pub fn push(&mut self, chunk: MessageChunk) {
match self { match self {
MessageContent::SingleText(text) => { MessageContent::SingleText(text) => {
*self = *self = MessageContent::MultipleChunks(vec![
MessageContent::MultipleChunks(vec![MessageChunk::Text { text: text.clone() }]); MessageChunk::Text { text: text.clone() },
chunk,
]);
} }
MessageContent::MultipleChunks(chunks) => { MessageContent::MultipleChunks(chunks) => {
chunks.push(chunk); chunks.push(chunk);
@ -1348,6 +1350,35 @@ mod tests {
); );
} }
#[test]
fn test_message_content_append() {
let mut content = MessageContent::SingleText("Initial text".to_string());
let chunk = MessageChunk::Text {
text: "Additional text".to_string(),
};
content.push(chunk);
match content {
MessageContent::MultipleChunks(chunks) => {
assert_eq!(chunks.len(), 2);
assert_eq!(
chunks[0],
MessageChunk::Text {
text: "Initial text".to_string()
}
);
assert_eq!(
chunks[1],
MessageChunk::Text {
text: "Additional text".to_string()
}
);
}
_ => panic!("Expected MultipleChunks, but got a different variant"),
}
}
#[test] #[test]
fn test_chat_request() { fn test_chat_request() {
let json = json!({ let json = json!({