[Bugfix][Rust] Fix HF multi-turn dataset integration (#56915)
Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> Co-authored-by: OpenAI Codex <codex@openai.com>
Kevin H. Luu committed
Sep 14, 2026 at 18:04 UTC
772e4f05fa929491f027e61a73876a95e3d86e64
2 files changed
+68
-10
rust/src/bench/src/datasets/multi_turn.rs
+57
@@ -350,6 +350,27 @@ pub fn load_sharegpt_multi_turn(
350
.as_array()
351
.ok_or_else(|| BenchError::Config("ShareGPT file must contain a JSON array".into()))?;
352
353
+ load_sharegpt_multi_turn_from_rows(
354
+ tokenizer,
355
+ entries,
356
+ num_conversations,
357
+ output_len_override,
358
+ max_turns,
359
+ seed,
360
+ request_id_prefix,
361
+ )
362
+}
363
+
364
+/// Load multi-turn conversations from deserialized ShareGPT rows.
365
+pub fn load_sharegpt_multi_turn_from_rows(
366
+ tokenizer: &TokenizerKind,
367
+ entries: &[serde_json::Value],
368
+ num_conversations: usize,
369
+ output_len_override: Option<usize>,
370
+ max_turns: Option<usize>,
371
+ seed: u64,
372
+ request_id_prefix: &str,
373
+) -> Result<Vec<MultiTurnConversation>> {
374
// Filter entries with at least 4 messages (2 turns: user+assistant+user+assistant)
375
let mut filtered: Vec<&serde_json::Value> = entries
376
.iter()
@@ -507,6 +528,42 @@ fn gen_prompt_to_target_len(
528
mod tests {
529
use super::*;
530
531
+ #[test]
532
+ fn test_load_sharegpt_multi_turn_from_hf_rows() {
533
+ let tokenizer =
534
+ TokenizerKind::Tiktoken(crate::tiktoken::load_builtin_tiktoken("gpt2").unwrap());
535
+ let rows = vec![serde_json::json!({
536
+ "conversations": [
537
+ {"from": "human", "value": "hello"},
538
+ {"from": "gpt", "value": "hi"},
539
+ {"from": "human", "value": "how are you?"},
540
+ {"from": "gpt", "value": "great"}
541
+ ]
542
+ })];
543
+
544
+ let conversations =
545
+ load_sharegpt_multi_turn_from_rows(&tokenizer, &rows, 1, Some(7), None, 42, "hf-")
546
+ .unwrap();
547
+ let actual = conversations
548
+ .iter()
549
+ .map(|conversation| {
550
+ (
551
+ conversation.conversation_id.as_str(),
552
+ conversation
553
+ .turns
554
+ .iter()
555
+ .map(|turn| (turn.user_message.as_ref(), turn.expected_output_len))
556
+ .collect::<Vec<_>>(),
557
+ )
558
+ })
559
+ .collect::<Vec<_>>();
560
+
561
+ assert_eq!(
562
+ actual,
563
+ vec![("hf-conv-0", vec![("hello", 7), ("how are you?", 7)])]
564
+ );
565
+ }
566
+
567
fn common_prefix_bytes(strings: &[&str]) -> usize {
568
if strings.is_empty() {
569
return 0;
rust/src/bench/src/multi_turn.rs
+11
-10
@@ -192,17 +192,18 @@ pub async fn run_multi_turn_benchmark(config: &BenchConfig) -> Result<serde_json
192
let dataset_id = config.dataset_path.as_deref().ok_or_else(|| {
193
BenchError::Config("--dataset-path is required for --dataset-name hf".into())
194
})?;
195
- let (downloaded_path, _config, _split) =
196
- crate::datasets::hf_dataset::download_hf_dataset(
197
- dataset_id,
198
- config.hf_subset.as_deref(),
199
- config.hf_split.as_deref(),
200
- config.num_prompts,
201
- )
202
- .await?;
203
- crate::datasets::multi_turn::load_sharegpt_multi_turn(
195
+ let (rows, _config, _split) = crate::datasets::hf_dataset::download_hf_dataset(
196
+ dataset_id,
197
+ config.hf_subset.as_deref(),
198
+ config.hf_split.as_deref(),
199
+ config.num_prompts,
200
+ config.seed,
201
+ config.disable_shuffle,
202
+ )
203
+ .await?;
204
+ crate::datasets::multi_turn::load_sharegpt_multi_turn_from_rows(
205
tok,
205
- &downloaded_path,
206
+ &rows,
207
config.num_prompts,
208
config.hf_output_len,
209
config.sharegpt_multi_turn_max_turns,