use anyhow::{Context, Result, bail}; use serde::Deserialize; use std::path::{Path, PathBuf}; #[derive(Debug, Clone, PartialEq, Eq)] pub struct AimeCase { pub id: String, pub problem: String, pub answer: String, } #[derive(Debug, Clone, PartialEq, Eq)] pub struct LoadedAimeCases { pub cases: Vec, pub local_path: PathBuf, } #[derive(Debug, Deserialize)] struct NormalizedAimeCase { id: String, problem: String, answer: String, } pub fn load_cases(data_dir: &Path) -> Result { let dataset_dir = data_dir.join("aime2026"); if !dataset_dir.exists() { bail!("missing local dataset; run lq_token_test dataset fetch aime2026"); } let jsonl_path = dataset_dir.join("aime2026.jsonl"); if jsonl_path.exists() { let contents = std::fs::read_to_string(&jsonl_path) .with_context(|| format!("failed to read {}", jsonl_path.display()))?; let cases = contents .lines() .enumerate() .filter(|(_, line)| !line.trim().is_empty()) .map(|(index, line)| { let parsed: NormalizedAimeCase = serde_json::from_str(line).with_context(|| { format!( "failed to parse normalized AIME JSONL line {} in {}", index + 1, jsonl_path.display() ) })?; Ok(AimeCase { id: parsed.id, problem: parsed.problem, answer: parsed.answer, }) }) .collect::>>()?; return Ok(LoadedAimeCases { cases, local_path: jsonl_path, }); } let has_parquet = std::fs::read_dir(&dataset_dir) .with_context(|| format!("failed to read {}", dataset_dir.display()))? .filter_map(Result::ok) .any(|entry| entry.path().extension().is_some_and(|ext| ext == "parquet")); if has_parquet { bail!( "AIME 2026 parquet conversion is not implemented yet; provide normalized aime2026.jsonl with fields id, problem, answer in {}", dataset_dir.display() ); } bail!("missing local dataset; run lq_token_test dataset fetch aime2026"); } impl AimeCase { pub fn prompt(&self) -> String { format!( "Solve the following AIME problem.\n\n{}\n\nRespond with the final integer answer.", self.problem ) } } #[cfg(test)] mod tests { use super::*; #[test] fn prompt_contains_problem_and_final_integer_instruction() { let case = AimeCase { id: "aime-1".to_string(), problem: "What is 20 + 22?".to_string(), answer: "42".to_string(), }; let prompt = case.prompt(); assert!(prompt.contains("What is 20 + 22?")); assert!(prompt.contains("final integer answer")); } #[test] fn loads_normalized_jsonl_cases() { let temp_dir = tempfile::tempdir().expect("create temp dir"); let dataset_dir = temp_dir.path().join("aime2026"); std::fs::create_dir_all(&dataset_dir).expect("create dataset dir"); std::fs::write( dataset_dir.join("aime2026.jsonl"), r#"{"id":"aime-1","problem":"What is 2+2?","answer":"4"} {"id":"aime-2","problem":"What is 3+5?","answer":"8"} "#, ) .expect("write jsonl"); let loaded = load_cases(temp_dir.path()).expect("load cases"); assert_eq!(loaded.local_path, dataset_dir.join("aime2026.jsonl")); assert_eq!(loaded.cases.len(), 2); assert_eq!(loaded.cases[0].answer, "4"); } #[test] fn loads_jsonl_normalized_from_hugging_face_rows() { let temp_dir = tempfile::tempdir().expect("create temp dir"); let dataset_dir = temp_dir.path().join("aime2026"); std::fs::create_dir_all(&dataset_dir).expect("create dataset dir"); std::fs::write( dataset_dir.join("aime2026.jsonl"), r#"{"id":"aime2026-7","problem":"Compute 6 times 7.","answer":"42"} "#, ) .expect("write normalized fetch output"); let loaded = load_cases(temp_dir.path()).expect("load cases"); assert_eq!(loaded.cases.len(), 1); assert_eq!(loaded.cases[0].id, "aime2026-7"); assert_eq!(loaded.cases[0].problem, "Compute 6 times 7."); assert_eq!(loaded.cases[0].answer, "42"); } #[test] fn parquet_only_returns_conversion_hint() { let temp_dir = tempfile::tempdir().expect("create temp dir"); let dataset_dir = temp_dir.path().join("aime2026"); std::fs::create_dir_all(&dataset_dir).expect("create dataset dir"); std::fs::write(dataset_dir.join("train.parquet"), b"parquet").expect("write parquet"); let error = load_cases(temp_dir.path()).expect_err("parquet should not load"); assert!( error .to_string() .contains("parquet conversion is not implemented yet") ); assert!(error.to_string().contains("aime2026.jsonl")); } }