miyagi 0.1.2

Sparse XOR adaptation and search for true binary GGUF language models
use std::fs;
use std::path::Path;

use regex::Regex;
use serde::{Deserialize, Serialize};
use serde_json::Value;

use crate::backend::{GenerateConfig, MiyagiBackend};
use crate::error::{Error, Result};

#[derive(Clone, Debug)]
pub struct DatasetConfig {
    pub question_field: String,
    pub answer_field: String,
    pub prompt_template: String,
    pub answer_regex: Regex,
    pub gold_regex: Regex,
    pub limit: Option<usize>,
    pub generation: GenerateConfig,
}

#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct DatasetCaseResult {
    pub index: usize,
    pub question: String,
    pub expected: String,
    pub predicted: Option<String>,
    pub correct: bool,
    pub response: String,
}

#[derive(Clone, Debug, Default, Deserialize, Serialize)]
pub struct DatasetReport {
    pub correct: usize,
    pub total: usize,
    pub cases: Vec<DatasetCaseResult>,
}

pub fn load_records(path: impl AsRef<Path>) -> Result<Vec<Value>> {
    let text = fs::read_to_string(path)?;
    if text.trim_start().starts_with('[') {
        return Ok(serde_json::from_str(&text)?);
    }
    text.lines()
        .filter(|line| !line.trim().is_empty())
        .map(|line| Ok(serde_json::from_str(line)?))
        .collect()
}

pub fn evaluate_dataset<B: MiyagiBackend>(
    backend: &mut B,
    records: &[Value],
    config: &DatasetConfig,
) -> Result<DatasetReport> {
    let mut report = DatasetReport::default();
    let limit = config.limit.unwrap_or(records.len()).min(records.len());
    for (index, record) in records.iter().take(limit).enumerate() {
        let question = string_field(record, &config.question_field, index)?;
        let raw_gold = string_field(record, &config.answer_field, index)?;
        let expected =
            capture(&config.gold_regex, raw_gold).ok_or_else(|| Error::InvalidDatasetRecord {
                index,
                reason: "gold answer regex did not match".to_owned(),
            })?;
        let prompt = config.prompt_template.replace("{question}", question);
        let response = backend.generate(&prompt, &config.generation)?;
        let predicted = capture(&config.answer_regex, &response);
        let correct = predicted.as_deref() == Some(expected.as_str());
        report.total += 1;
        report.correct += usize::from(correct);
        report.cases.push(DatasetCaseResult {
            index,
            question: question.to_owned(),
            expected,
            predicted,
            correct,
            response,
        });
    }
    Ok(report)
}

fn string_field<'a>(record: &'a Value, field: &str, index: usize) -> Result<&'a str> {
    record
        .get(field)
        .and_then(Value::as_str)
        .ok_or_else(|| Error::InvalidDatasetRecord {
            index,
            reason: format!("field {field:?} is missing or is not a string"),
        })
}

fn capture(regex: &Regex, text: &str) -> Option<String> {
    regex
        .captures(text)
        .and_then(|captures| captures.get(1))
        .map(|capture| capture.as_str().replace(',', ""))
}