|
- 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<AimeCase>,
- pub local_path: PathBuf,
- }
-
- #[derive(Debug, Deserialize)]
- struct NormalizedAimeCase {
- id: String,
- problem: String,
- answer: String,
- }
-
- pub fn load_cases(data_dir: &Path) -> Result<LoadedAimeCases> {
- 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::<Result<Vec<_>>>()?;
- 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"));
- }
- }
|