Unverified Commit c9fad85c authored by Ayush Agarwal's avatar Ayush Agarwal Committed by GitHub
Browse files

chore: jamba based parsers (#4776)


Signed-off-by: default avatarayushag <ayushag@nvidia.com>
parent d2c23e41
...@@ -38,6 +38,7 @@ Parser to Model Mapping ...@@ -38,6 +38,7 @@ Parser to Model Mapping
| deepseek_v3 | deepseek-ai/DeepSeek-V3, deepseek-ai/DeepSeek-R1, deepseek-ai/DeepSeek-R1-0528 | | deepseek_v3 | deepseek-ai/DeepSeek-V3, deepseek-ai/DeepSeek-R1, deepseek-ai/DeepSeek-R1-0528 |
| deepseek_v3_1 | deepseek-ai/DeepSeek-V3.1 | | deepseek_v3_1 | deepseek-ai/DeepSeek-V3.1 |
| pythonic | meta-llama/Llama-4-* | | pythonic | meta-llama/Llama-4-* |
| jamba | ai21labs/AI21-Jamba-*-1.5, ai21labs/AI21-Jamba-*-1.6, ai21labs/AI21-Jamba-*-1.7, |
## Examples ## Examples
......
...@@ -239,4 +239,14 @@ impl ToolCallConfig { ...@@ -239,4 +239,14 @@ impl ToolCallConfig {
parser_config: ParserConfig::Xml(XmlParserConfig::default()), parser_config: ParserConfig::Xml(XmlParserConfig::default()),
} }
} }
pub fn jamba() -> Self {
Self {
parser_config: ParserConfig::Json(JsonParserConfig {
tool_call_start_tokens: vec!["<tool_calls>".to_string()],
tool_call_end_tokens: vec!["</tool_calls>".to_string()],
..Default::default()
}),
}
}
} }
...@@ -36,6 +36,7 @@ pub fn get_tool_parser_map() -> &'static HashMap<&'static str, ToolCallConfig> { ...@@ -36,6 +36,7 @@ pub fn get_tool_parser_map() -> &'static HashMap<&'static str, ToolCallConfig> {
map.insert("deepseek_v3", ToolCallConfig::deepseek_v3()); map.insert("deepseek_v3", ToolCallConfig::deepseek_v3());
map.insert("deepseek_v3_1", ToolCallConfig::deepseek_v3_1()); map.insert("deepseek_v3_1", ToolCallConfig::deepseek_v3_1());
map.insert("qwen3_coder", ToolCallConfig::qwen3_coder()); map.insert("qwen3_coder", ToolCallConfig::qwen3_coder());
map.insert("jamba", ToolCallConfig::jamba());
map.insert("default", ToolCallConfig::default()); map.insert("default", ToolCallConfig::default());
map map
}) })
...@@ -191,6 +192,7 @@ mod tests { ...@@ -191,6 +192,7 @@ mod tests {
"deepseek_v3", "deepseek_v3",
"deepseek_v3_1", "deepseek_v3_1",
"qwen3_coder", "qwen3_coder",
"jamba",
]; ];
for parser in available_parsers { for parser in available_parsers {
assert!(parsers.contains(&parser)); assert!(parsers.contains(&parser));
...@@ -940,19 +942,11 @@ Remember, San Francisco weather can be quite unpredictable, particularly with it ...@@ -940,19 +942,11 @@ Remember, San Francisco weather can be quite unpredictable, particularly with it
} }
#[tokio::test] #[tokio::test]
#[ignore]
async fn test_ai21labs_ai21_jamba_15_mini_simple() { async fn test_ai21labs_ai21_jamba_15_mini_simple() {
let input = r#" [ let input = r#"<tool_calls>[
{"name": "get_weather", "arguments": {"location": "San Francisco, CA", "unit": "fahrenheit"}} {"name": "get_weather", "arguments": {"location": "San Francisco, CA", "unit": "fahrenheit"}}
]"#; ]</tool_calls>"#;
let config = ToolCallConfig { let config = ToolCallConfig::jamba();
parser_config: ParserConfig::Json(JsonParserConfig {
tool_call_start_tokens: vec![],
tool_call_end_tokens: vec![],
arguments_keys: vec!["arguments".to_string()],
..Default::default()
}),
};
let (result, content) = try_tool_call_parse(input, &config).await.unwrap(); let (result, content) = try_tool_call_parse(input, &config).await.unwrap();
assert_eq!(content, Some("".to_string())); assert_eq!(content, Some("".to_string()));
assert!(!result.is_empty()); assert!(!result.is_empty());
...@@ -963,6 +957,29 @@ Remember, San Francisco weather can be quite unpredictable, particularly with it ...@@ -963,6 +957,29 @@ Remember, San Francisco weather can be quite unpredictable, particularly with it
assert_eq!(args["unit"], "fahrenheit"); assert_eq!(args["unit"], "fahrenheit");
} }
#[tokio::test]
async fn test_ai21labs_ai21_jamba_15_mini_multiple() {
let input = r#"<tool_calls>[
{"name": "get_weather", "arguments": {"location": "San Francisco, CA", "unit": "fahrenheit"}},
{"name": "get_weather", "arguments": {"location": "New York, NY", "unit": "celsius"}}
]</tool_calls>"#;
let config = ToolCallConfig::jamba();
let (result, content) = try_tool_call_parse(input, &config).await.unwrap();
assert_eq!(content, Some("".to_string()));
assert!(!result.is_empty());
assert_eq!(result.len(), 2);
let (name, args) = extract_name_and_args(result[0].clone());
assert_eq!(name, "get_weather");
assert_eq!(args["location"], "San Francisco, CA");
assert_eq!(args["unit"], "fahrenheit");
let (name, args) = extract_name_and_args(result[1].clone());
assert_eq!(name, "get_weather");
assert_eq!(args["location"], "New York, NY");
assert_eq!(args["unit"], "celsius");
}
#[tokio::test] #[tokio::test]
#[ignore] #[ignore]
async fn test_salesforce_llama_xlam_2_8b_fc_r_simple() { async fn test_salesforce_llama_xlam_2_8b_fc_r_simple() {
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment