captchaforge 0.2.36

//! VLM training dataset format for captcha-specific fine-tuning.
//!
//! Vision-language captcha solving currently uses general-purpose
//! models (Ollama llama3.2-vision, GPT-4V, Claude vision) which
//! score ~70% on bench image-grid suites. A captcha-specific
//! fine-tune of a 7B-parameter VLM should hit 90%+ — the
//! bottleneck is collecting + structuring the training data.
//!
//! [`VlmDataset`] is the storage substrate: per-sample
//! (screenshot + prompt + ground-truth answer) tuples written to a
//! manifest JSON pointing at PNGs on disk. The format is
//! deliberately simple + human-auditable so the manifest can be
//! reviewed pre-training without a Python harness.
//!
//! ## Three sample sources
//!
//! 1. **Bench fixtures** — every fixture in
//!    `bench/src/fixtures.rs` has a known correct answer. Loop
//!    them, screenshot each, write a sample. Free, deterministic.
//! 2. **Failure corpus** — every solver failure that captures a
//!    screenshot gets re-labelled by an operator (or by GPT-4V as
//!    a labelling oracle) and added to the training set. Most
//!    valuable signal per sample because failures are where the
//!    model has the most to learn.
//! 3. **Synthetic generators** — programmatically generate image
//!    grids with known answers (e.g. "render N traffic lights at
//!    random positions, mark the cells containing them"). Cheap;
//!    high volume; lower per-sample value than (2) but useful for
//!    bootstrapping.
//!
//! ## Format
//!
//! - `manifest.json` — array of [`VlmSample`] entries.
//! - `images/<hash>.png` — screenshot per sample, content-addressed
//!   by SHA-256 of the image bytes so duplicate screenshots
//!   dedupe to a single file.
//!
//! Hand the directory to a training pipeline (`peft` / `trl` /
//! `axolotl`); the manifest schema matches what the most common
//! VLM fine-tune toolchains consume.

#![allow(dead_code)] // module is opt-in.

use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::path::{Path, PathBuf};

/// One training sample.
///
/// `prompt` is what we'd ask the model at inference time — what the
/// captcha widget's UI text says (e.g. "select all squares with
/// traffic lights"). `answer` is the structured ground truth — for
/// click captchas, the cell indices to click; for text captchas,
/// the transcription.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct VlmSample {
    /// Stable ID — content-addressed by image hash. Useful for
    /// "did we already train on this exact screenshot" dedupe.
    pub id: String,
    /// Vendor / captcha-type label (`"recaptcha-v2-image-grid"`,
    /// `"hcaptcha-icon-pick"`, `"funcaptcha-rotate"`).
    pub kind: String,
    /// Path to the PNG under the dataset's `images/` dir.
    pub image_path: String,
    /// Visible prompt text from the captcha widget.
    pub prompt: String,
    /// Structured answer — JSON string. Format depends on
    /// captcha kind:
    /// - image-grid: `{"cells": [0, 4, 7]}` (cell indices)
    /// - text: `{"transcription": "abcd1234"}`
    /// - click: `{"x": 240, "y": 180}` (viewport coords)
    /// - rotation: `{"degrees": 90}`
    pub answer: String,
    /// Optional negative-example marker — same image, wrong answer
    /// the model should LEARN to reject. Boosts robustness.
    pub is_negative: bool,
}

/// One on-disk VLM training dataset.
pub struct VlmDataset {
    root: PathBuf,
}

impl VlmDataset {
    /// Open or create a dataset at `path`. Creates `images/` if
    /// missing.
    pub fn open(path: impl AsRef<Path>) -> Result<Self> {
        let root = path.as_ref().to_path_buf();
        std::fs::create_dir_all(root.join("images"))
            .with_context(|| format!("creating dataset images dir under {}", root.display()))?;
        Ok(Self { root })
    }

    /// Borrow the dataset root.
    pub fn root(&self) -> &Path {
        &self.root
    }

    /// Add a sample. Writes the image to `images/<hash>.png`
    /// (deduping if an image with the same content hash is already
    /// present), updates the manifest.
    pub fn add_sample(
        &self,
        kind: impl Into<String>,
        prompt: impl Into<String>,
        answer: impl Into<String>,
        image_png: &[u8],
        is_negative: bool,
    ) -> Result<VlmSample> {
        // Content-address the image so identical screenshots dedupe.
        let id = sha256_hex(image_png);
        let image_filename = format!("{id}.png");
        let image_path = self.root.join("images").join(&image_filename);
        if !image_path.exists() {
            std::fs::write(&image_path, image_png)
                .with_context(|| format!("writing image {}", image_path.display()))?;
        }
        let sample = VlmSample {
            id: id.clone(),
            kind: kind.into(),
            image_path: format!("images/{image_filename}"),
            prompt: prompt.into(),
            answer: answer.into(),
            is_negative,
        };
        self.append_to_manifest(&sample)?;
        Ok(sample)
    }

    /// Load every sample from the manifest. Cheap for typical
    /// dataset sizes (<100k samples).
    pub fn load_all(&self) -> Result<Vec<VlmSample>> {
        let path = self.manifest_path();
        if !path.exists() {
            return Ok(Vec::new());
        }
        let raw = std::fs::read_to_string(&path)
            .with_context(|| format!("reading {}", path.display()))?;
        let mut out = Vec::new();
        for (i, line) in raw.lines().enumerate() {
            if line.trim().is_empty() {
                continue;
            }
            let sample: VlmSample = serde_json::from_str(line).with_context(|| {
                format!("parsing manifest line {} of {}", i + 1, path.display())
            })?;
            out.push(sample);
        }
        Ok(out)
    }

    /// Aggregate counts per kind — useful for "is the per-kind
    /// distribution balanced enough to train on?".
    pub fn kind_counts(&self) -> Result<Vec<(String, usize)>> {
        let samples = self.load_all()?;
        let mut counts: std::collections::HashMap<String, usize> = std::collections::HashMap::new();
        for s in samples {
            *counts.entry(s.kind).or_insert(0) += 1;
        }
        let mut out: Vec<(String, usize)> = counts.into_iter().collect();
        out.sort();
        Ok(out)
    }

    /// Render the dataset as a JSONL manifest path the caller can
    /// pass to a Python training harness.
    pub fn manifest_path(&self) -> PathBuf {
        self.root.join("manifest.jsonl")
    }

    fn append_to_manifest(&self, sample: &VlmSample) -> Result<()> {
        use std::io::Write;
        let mut line = serde_json::to_string(sample).context("serialising VLM sample")?;
        line.push('\n');
        let mut file = std::fs::OpenOptions::new()
            .create(true)
            .append(true)
            .open(self.manifest_path())
            .with_context(|| format!("opening manifest {}", self.manifest_path().display()))?;
        file.write_all(line.as_bytes())
            .context("writing manifest line")?;
        Ok(())
    }
}

fn sha256_hex(bytes: &[u8]) -> String {
    let mut hasher = Sha256::new();
    hasher.update(bytes);
    let digest = hasher.finalize();
    digest.iter().map(|b| format!("{b:02x}")).collect()
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::tempdir;

    #[test]
    fn open_creates_images_dir() {
        let tmp = tempdir().unwrap();
        let _ds = VlmDataset::open(tmp.path()).unwrap();
        assert!(tmp.path().join("images").exists());
    }

    #[test]
    fn add_sample_writes_image_and_manifest_line() {
        let tmp = tempdir().unwrap();
        let ds = VlmDataset::open(tmp.path()).unwrap();
        let sample = ds
            .add_sample(
                "recaptcha-v2-image-grid",
                "select all squares with traffic lights",
                r#"{"cells":[0,4,8]}"#,
                b"PNGFAKEBYTES",
                false,
            )
            .unwrap();
        assert!(tmp.path().join(&sample.image_path).exists());
        assert!(ds.manifest_path().exists());
        let loaded = ds.load_all().unwrap();
        assert_eq!(loaded.len(), 1);
        assert_eq!(loaded[0], sample);
    }

    #[test]
    fn add_sample_dedupes_identical_image_to_single_file() {
        let tmp = tempdir().unwrap();
        let ds = VlmDataset::open(tmp.path()).unwrap();
        let png = b"DUPLICATE_PNG_BYTES";
        let s1 = ds
            .add_sample("kind1", "prompt1", "answer1", png, false)
            .unwrap();
        let s2 = ds
            .add_sample("kind2", "prompt2", "answer2", png, false)
            .unwrap();
        // Both samples share the same content-addressed image path.
        assert_eq!(s1.image_path, s2.image_path);
        // But the manifest has BOTH samples (different prompts).
        let loaded = ds.load_all().unwrap();
        assert_eq!(loaded.len(), 2);
        // Image dir contains exactly one file.
        let images: Vec<_> = std::fs::read_dir(tmp.path().join("images"))
            .unwrap()
            .filter_map(|e| e.ok())
            .collect();
        assert_eq!(images.len(), 1);
    }

    #[test]
    fn kind_counts_aggregates_correctly() {
        let tmp = tempdir().unwrap();
        let ds = VlmDataset::open(tmp.path()).unwrap();
        for i in 0..5 {
            ds.add_sample(
                "image-grid",
                format!("prompt {i}"),
                "answer",
                format!("PNG_GRID_{i}").as_bytes(),
                false,
            )
            .unwrap();
        }
        for i in 0..3 {
            ds.add_sample(
                "icon-pick",
                format!("prompt {i}"),
                "answer",
                format!("PNG_ICON_{i}").as_bytes(),
                false,
            )
            .unwrap();
        }
        let counts = ds.kind_counts().unwrap();
        assert_eq!(counts.len(), 2);
        let map: std::collections::HashMap<_, _> = counts.into_iter().collect();
        assert_eq!(map.get("image-grid"), Some(&5));
        assert_eq!(map.get("icon-pick"), Some(&3));
    }

    #[test]
    fn negative_samples_round_trip_with_marker() {
        let tmp = tempdir().unwrap();
        let ds = VlmDataset::open(tmp.path()).unwrap();
        ds.add_sample("k", "p", "wrong-answer", b"X", true).unwrap();
        let loaded = ds.load_all().unwrap();
        assert_eq!(loaded.len(), 1);
        assert!(loaded[0].is_negative);
    }

    #[test]
    fn load_all_returns_empty_for_fresh_dataset() {
        let tmp = tempdir().unwrap();
        let ds = VlmDataset::open(tmp.path()).unwrap();
        assert!(ds.load_all().unwrap().is_empty());
    }

    #[test]
    fn sha256_hex_matches_known_oracle() {
        // SHA-256 of the empty string per RFC 6234.
        assert_eq!(
            sha256_hex(b""),
            "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"
        );
        // SHA-256("abc")
        assert_eq!(
            sha256_hex(b"abc"),
            "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad"
        );
    }
}