use anyhow::{Context, Result, bail}; use std::collections::HashMap; use std::path::{Path, PathBuf}; #[derive(Debug, Clone, PartialEq, Eq)] pub struct GpqaCase { pub id: String, pub question: String, pub choices: [(String, String); 4], pub answer: char, } #[derive(Debug, Clone, PartialEq, Eq)] pub struct LoadedGpqaCases { pub cases: Vec, pub local_path: PathBuf, } pub fn load_cases(data_dir: &Path) -> Result { let dataset_dir = data_dir.join("gpqa_diamond"); let csv_path = dataset_dir.join("gpqa_diamond.csv"); if !csv_path.exists() { bail!("missing local dataset; run lq_token_test dataset fetch gpqa-diamond"); } let contents = std::fs::read_to_string(&csv_path) .with_context(|| format!("failed to read {}", csv_path.display()))?; let rows = parse_csv(&contents).context("failed to parse GPQA CSV")?; let (headers, records) = rows.split_first().context("GPQA CSV is empty")?; let header_index = headers .iter() .enumerate() .map(|(index, header)| (normalize_header(header), index)) .collect::>(); let cases = records .iter() .enumerate() .map(|(index, record)| { let question = csv_value(record, &header_index, &["question"]) .with_context(|| format!("missing Question column in row {}", index + 2))?; let correct = csv_value(record, &header_index, &["correctanswer", "correct"]) .with_context(|| format!("missing Correct Answer column in row {}", index + 2))?; let incorrect_1 = csv_value(record, &header_index, &["incorrectanswer1", "incorrect1"]) .with_context(|| { format!("missing Incorrect Answer 1 column in row {}", index + 2) })?; let incorrect_2 = csv_value(record, &header_index, &["incorrectanswer2", "incorrect2"]) .with_context(|| { format!("missing Incorrect Answer 2 column in row {}", index + 2) })?; let incorrect_3 = csv_value(record, &header_index, &["incorrectanswer3", "incorrect3"]) .with_context(|| { format!("missing Incorrect Answer 3 column in row {}", index + 2) })?; let id = format!("gpqa-diamond-{}", index + 1); let (choices, answer) = rotated_choices( [ correct.to_string(), incorrect_1.to_string(), incorrect_2.to_string(), incorrect_3.to_string(), ], index % 4, ); Ok(GpqaCase { id, question: question.to_string(), choices, answer, }) }) .collect::>>()?; Ok(LoadedGpqaCases { cases, local_path: csv_path, }) } fn csv_value<'a>( record: &'a [String], header_index: &HashMap, names: &[&str], ) -> Option<&'a str> { names .iter() .find_map(|name| header_index.get(*name)) .and_then(|index| record.get(*index)) .map(String::as_str) .filter(|value| !value.trim().is_empty()) } fn normalize_header(header: &str) -> String { header .chars() .filter(|ch| ch.is_ascii_alphanumeric()) .collect::() .to_ascii_lowercase() } fn rotated_choices(mut values: [String; 4], rotation: usize) -> ([(String, String); 4], char) { let rotation = rotation % values.len(); values.rotate_left(rotation); let answer_index = (4 - rotation) % 4; let labels = ["A", "B", "C", "D"]; let answer = labels[answer_index] .chars() .next() .expect("labels are non-empty"); let choices = std::array::from_fn(|index| (labels[index].to_string(), values[index].clone())); (choices, answer) } fn parse_csv(input: &str) -> Result>> { let mut rows = Vec::new(); let mut row = Vec::new(); let mut field = String::new(); let mut chars = input.chars().peekable(); let mut in_quotes = false; while let Some(ch) = chars.next() { match ch { '"' if in_quotes && chars.peek() == Some(&'"') => { field.push('"'); chars.next(); } '"' => in_quotes = !in_quotes, ',' if !in_quotes => { row.push(std::mem::take(&mut field)); } '\n' if !in_quotes => { row.push(std::mem::take(&mut field)); if !row.iter().all(|value| value.is_empty()) { rows.push(std::mem::take(&mut row)); } else { row.clear(); } } '\r' if !in_quotes => {} _ => field.push(ch), } } if in_quotes { bail!("unterminated quoted field"); } if !field.is_empty() || !row.is_empty() { row.push(field); rows.push(row); } Ok(rows) } impl GpqaCase { pub fn prompt(&self) -> String { let choices = self .choices .iter() .map(|(label, text)| format!("{label}. {text}")) .collect::>() .join("\n"); format!( "Answer the following multiple-choice question.\n\n{}\n\n{}\n\nPlease answer with exactly one letter: A, B, C, or D.", self.question, choices ) } } #[cfg(test)] mod tests { use super::*; #[test] fn prompt_contains_choices_and_single_letter_instruction() { let case = GpqaCase { id: "gpqa-1".to_string(), question: "Which option is correct?".to_string(), choices: [ ("A".to_string(), "Alpha".to_string()), ("B".to_string(), "Beta".to_string()), ("C".to_string(), "Gamma".to_string()), ("D".to_string(), "Delta".to_string()), ], answer: 'C', }; let prompt = case.prompt(); assert!(prompt.contains("Which option is correct?")); assert!(prompt.contains("A. Alpha")); assert!(prompt.contains("B. Beta")); assert!(prompt.contains("C. Gamma")); assert!(prompt.contains("D. Delta")); assert!(prompt.contains("answer with exactly one letter")); } #[test] fn loads_gpqa_csv_with_common_columns() { let temp_dir = tempfile::tempdir().expect("create temp dir"); let dataset_dir = temp_dir.path().join("gpqa_diamond"); std::fs::create_dir_all(&dataset_dir).expect("create dataset dir"); std::fs::write( dataset_dir.join("gpqa_diamond.csv"), "Question,Correct Answer,Incorrect Answer 1,Incorrect Answer 2,Incorrect Answer 3\n\"Which is correct?\",\"Right\",\"Wrong, with comma\",\"Nope\",\"Never\"\n\"Which is second?\",\"Second right\",\"Second wrong 1\",\"Second wrong 2\",\"Second wrong 3\"\n", ) .expect("write csv"); let loaded = load_cases(temp_dir.path()).expect("load cases"); assert_eq!(loaded.local_path, dataset_dir.join("gpqa_diamond.csv")); assert_eq!(loaded.cases.len(), 2); assert_eq!(loaded.cases[0].question, "Which is correct?"); assert_eq!( loaded.cases[0].choices[0], ("A".to_string(), "Right".to_string()) ); assert_eq!(loaded.cases[0].answer, 'A'); assert_eq!(loaded.cases[1].answer, 'D'); assert_eq!( loaded.cases[1] .choices .iter() .find(|(label, _)| label == "D") .map(|(_, text)| text.as_str()), Some("Second right") ); assert!(loaded.cases.iter().any(|case| case.answer != 'A')); } }