Spaces:
Paused
Paused
File size: 3,041 Bytes
2f86d9b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 | # 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.
|