|
- 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<GpqaCase>,
- pub local_path: PathBuf,
- }
-
- pub fn load_cases(data_dir: &Path) -> Result<LoadedGpqaCases> {
- 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::<HashMap<_, _>>();
-
- 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::<Result<Vec<_>>>()?;
-
- Ok(LoadedGpqaCases {
- cases,
- local_path: csv_path,
- })
- }
-
- fn csv_value<'a>(
- record: &'a [String],
- header_index: &HashMap<String, usize>,
- 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::<String>()
- .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<Vec<Vec<String>>> {
- 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::<Vec<_>>()
- .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'));
- }
- }
|