api模型检测
Вы не можете выбрать более 25 тем Темы должны начинаться с буквы или цифры, могут содержать дефисы(-) и должны содержать не более 35 символов.

158 строки
5.1 KiB

  1. use anyhow::{Context, Result, bail};
  2. use serde::Deserialize;
  3. use std::path::{Path, PathBuf};
  4. #[derive(Debug, Clone, PartialEq, Eq)]
  5. pub struct AimeCase {
  6. pub id: String,
  7. pub problem: String,
  8. pub answer: String,
  9. }
  10. #[derive(Debug, Clone, PartialEq, Eq)]
  11. pub struct LoadedAimeCases {
  12. pub cases: Vec<AimeCase>,
  13. pub local_path: PathBuf,
  14. }
  15. #[derive(Debug, Deserialize)]
  16. struct NormalizedAimeCase {
  17. id: String,
  18. problem: String,
  19. answer: String,
  20. }
  21. pub fn load_cases(data_dir: &Path) -> Result<LoadedAimeCases> {
  22. let dataset_dir = data_dir.join("aime2026");
  23. if !dataset_dir.exists() {
  24. bail!("missing local dataset; run lq_token_test dataset fetch aime2026");
  25. }
  26. let jsonl_path = dataset_dir.join("aime2026.jsonl");
  27. if jsonl_path.exists() {
  28. let contents = std::fs::read_to_string(&jsonl_path)
  29. .with_context(|| format!("failed to read {}", jsonl_path.display()))?;
  30. let cases = contents
  31. .lines()
  32. .enumerate()
  33. .filter(|(_, line)| !line.trim().is_empty())
  34. .map(|(index, line)| {
  35. let parsed: NormalizedAimeCase = serde_json::from_str(line).with_context(|| {
  36. format!(
  37. "failed to parse normalized AIME JSONL line {} in {}",
  38. index + 1,
  39. jsonl_path.display()
  40. )
  41. })?;
  42. Ok(AimeCase {
  43. id: parsed.id,
  44. problem: parsed.problem,
  45. answer: parsed.answer,
  46. })
  47. })
  48. .collect::<Result<Vec<_>>>()?;
  49. return Ok(LoadedAimeCases {
  50. cases,
  51. local_path: jsonl_path,
  52. });
  53. }
  54. let has_parquet = std::fs::read_dir(&dataset_dir)
  55. .with_context(|| format!("failed to read {}", dataset_dir.display()))?
  56. .filter_map(Result::ok)
  57. .any(|entry| entry.path().extension().is_some_and(|ext| ext == "parquet"));
  58. if has_parquet {
  59. bail!(
  60. "AIME 2026 parquet conversion is not implemented yet; provide normalized aime2026.jsonl with fields id, problem, answer in {}",
  61. dataset_dir.display()
  62. );
  63. }
  64. bail!("missing local dataset; run lq_token_test dataset fetch aime2026");
  65. }
  66. impl AimeCase {
  67. pub fn prompt(&self) -> String {
  68. format!(
  69. "Solve the following AIME problem.\n\n{}\n\nRespond with the final integer answer.",
  70. self.problem
  71. )
  72. }
  73. }
  74. #[cfg(test)]
  75. mod tests {
  76. use super::*;
  77. #[test]
  78. fn prompt_contains_problem_and_final_integer_instruction() {
  79. let case = AimeCase {
  80. id: "aime-1".to_string(),
  81. problem: "What is 20 + 22?".to_string(),
  82. answer: "42".to_string(),
  83. };
  84. let prompt = case.prompt();
  85. assert!(prompt.contains("What is 20 + 22?"));
  86. assert!(prompt.contains("final integer answer"));
  87. }
  88. #[test]
  89. fn loads_normalized_jsonl_cases() {
  90. let temp_dir = tempfile::tempdir().expect("create temp dir");
  91. let dataset_dir = temp_dir.path().join("aime2026");
  92. std::fs::create_dir_all(&dataset_dir).expect("create dataset dir");
  93. std::fs::write(
  94. dataset_dir.join("aime2026.jsonl"),
  95. r#"{"id":"aime-1","problem":"What is 2+2?","answer":"4"}
  96. {"id":"aime-2","problem":"What is 3+5?","answer":"8"}
  97. "#,
  98. )
  99. .expect("write jsonl");
  100. let loaded = load_cases(temp_dir.path()).expect("load cases");
  101. assert_eq!(loaded.local_path, dataset_dir.join("aime2026.jsonl"));
  102. assert_eq!(loaded.cases.len(), 2);
  103. assert_eq!(loaded.cases[0].answer, "4");
  104. }
  105. #[test]
  106. fn loads_jsonl_normalized_from_hugging_face_rows() {
  107. let temp_dir = tempfile::tempdir().expect("create temp dir");
  108. let dataset_dir = temp_dir.path().join("aime2026");
  109. std::fs::create_dir_all(&dataset_dir).expect("create dataset dir");
  110. std::fs::write(
  111. dataset_dir.join("aime2026.jsonl"),
  112. r#"{"id":"aime2026-7","problem":"Compute 6 times 7.","answer":"42"}
  113. "#,
  114. )
  115. .expect("write normalized fetch output");
  116. let loaded = load_cases(temp_dir.path()).expect("load cases");
  117. assert_eq!(loaded.cases.len(), 1);
  118. assert_eq!(loaded.cases[0].id, "aime2026-7");
  119. assert_eq!(loaded.cases[0].problem, "Compute 6 times 7.");
  120. assert_eq!(loaded.cases[0].answer, "42");
  121. }
  122. #[test]
  123. fn parquet_only_returns_conversion_hint() {
  124. let temp_dir = tempfile::tempdir().expect("create temp dir");
  125. let dataset_dir = temp_dir.path().join("aime2026");
  126. std::fs::create_dir_all(&dataset_dir).expect("create dataset dir");
  127. std::fs::write(dataset_dir.join("train.parquet"), b"parquet").expect("write parquet");
  128. let error = load_cases(temp_dir.path()).expect_err("parquet should not load");
  129. assert!(
  130. error
  131. .to_string()
  132. .contains("parquet conversion is not implemented yet")
  133. );
  134. assert!(error.to_string().contains("aime2026.jsonl"));
  135. }
  136. }