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:
parent
19ea85f8dc
commit
9a7830bd28
|
@ -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);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
|
@ -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!({
|
||||||
|
|
Loading…
Reference in New Issue