gemma/diffusion/hackable_diffusion_adapter/eval/README.md

Evaluation Metrics and Helpers for HD Text

This package contains all evaluation pipelines, custom metrics, and helpers used during the SFT training of Hackable Diffusion (HD) text models.

Using custom JAX-compatible and host-side metrics, it tracks performance across different SFT tasks (Sudoku, PubMedQA).

Key Components

1. Base Metric Classes (base_metric.py)

Shared base classes (BaseTokenizerMetric, BaseTextMetric, BaseSimpleTextMetric) that consolidate tokenizer access, 3D→2D squeezing, and host-side io_callback wiring into a reusable hierarchy.

2. PubMedQA Evaluation (pubmedqa_eval.py)

Contains evaluation helpers for the PubMedQA dataset (both short yes/no/maybe answers and long explanations).

• extract_pubmedqa_answer: Extracts the structured answer from the model's response by looking for the "The answer is: {yes|no|maybe}" marker.
• PubMedQAAccuracy: Computes accuracy of the extracted answers against the ground truth.
• BLEUScore: Computes smoothed sentence-level BLEU between generated and ground-truth text using sacrebleu.

3. Sudoku Evaluation (sudoku_eval.py)

A comprehensive evaluation suite for Sudoku puzzle solving.

• SudokuAllMetrics: A unified host-side metric that computes 20 different stratified accuracy metrics (e.g., overall cell accuracy, puzzle-level correctness, accuracy stratified by difficulty (easy/medium/hard), and count of completely correct puzzles). Grouping these under a single Kauldron metric key avoids TPU HBM OOM.

• Hardness Stratification: Difficulty is categorized dynamically on-device according to the number of non-zero clue digits in the initial puzzle context:

  • All: Every puzzle in the eval dataset.
  • Easy: Puzzles with $\ge 40$ clues.
  • Medium: Puzzles with $\ge 30$ and $< 40$ clues.
  • Hard: Puzzles with $< 30$ clues.

• Monitored Metrics Modes:

  • accuracy: Exact grid match correctness (requires all 81 digits to match reference).
  • cell_accuracy: Fraction of correctly infilled digits.
  • exact_mask_accuracy: Exact preservation correctness of all initial clues.
  • partial_mask_accuracy: Fraction of initial clues preserved.
  • count: Tracks the quantity of evaluated instances in each difficulty category.

4. Text Metrics (text_metric.py)

A utility metric to decode and log generated samples alongside prompts during training.

• DetokenizePromptAndResponse: A custom Kauldron metric that collects prompt and response tokens, decodes them using the SentencePiece tokenizer, and logs the detokenized text strings.

5. Autoregressive Evaluation (ar_eval.py)

Contains autoregressive evaluation helpers for DiffusionGemma SFT.

• make_ar_evals: Creates autoregressive evaluation configurations.
• AR_DENOISING_STEPS: Defines the denoising step counts for AR evaluation.