# Repetition Diagnostic

A builder tool for tiny language model developers. Sweeps `repetition_penalty` values and measures how they affect generation quality.

## Why

Tiny models (<50M params) are extremely prone to degenerate repetition loops. The default `repetition_penalty=1.0` means **no penalty** — the model can output the same token forever. Most training scripts don't test generation quality, so you discover this at inference time with no idea what to do.

This tool gives you a quick diagnostic: sweep the penalty, see where quality peaks, and know what to set in your inference config.

## What it measures

| Metric | Description |
|--------|-------------|
| **Repetition rate** | Fraction of output tokens that already appeared earlier in the sequence |
| **Unique token ratio** | Unique tokens / total tokens (diversity) |
| **Max loop length** | Longest repeated n-gram (how deep the loop goes) |
| **Repeated trigrams** | Count of 3-token sequences appearing 2+ times |

## Usage

```bash
# Basic: sweep penalties on a single prompt
python3 repetition_diagnostic.py --model gpt2 --prompt "Once upon a time"

# Multiple prompts
python3 repetition_diagnostic.py --model gpt2 --prompts "Once upon a time|The cat sat|I want to"

# Custom penalty range
python3 repetition_diagnostic.py --model gpt2 --penalties "1.0,1.05,1.1,1.2,1.5"

# Save results as JSON
python3 repetition_diagnostic.py --model gpt2 --output results.json
```

## Example output

```
Model: gpt2
Prompts: ['Once upon a time']
Penalties: [1.0, 1.1, 1.2, 1.5, 2.0]

Prompt                   RP Tokens   Repet  Unique  MaxLoop  RepTri
------------------------------------------------------------------------
Once upon a time        1.0     48   25.5%   72.9%        3       1
Once upon a time        1.1     48    0.0%  100.0%        0       0
Once upon a time        1.2     48    0.0%  100.0%        0       0
Once upon a time        1.5     48    0.0%  100.0%        0       0
Once upon a time        2.0     48    0.0%  100.0%        0       0
```

## Interpretation

- **rp=1.0 with high repetition rate**: Your model has a repetition problem. Try rp=1.1 or 1.2.
- **rp=1.1 already fixes it**: Good — your model is mostly fine, just needs a small nudge.
- **Even rp=2.0 still shows loops**: The model is severely undertrained or has a degenerate mode. No penalty will save it.
- **rp=1.0 shows no repetition**: Your model is robust. You can use the default.

## Requirements

- Python 3.8+
- `torch`
- `transformers`

```bash
pip install torch transformers
```

## Notes

- Uses the same seed for all penalty values so the comparison is fair (same random draws, different penalty).
- Generation is done manually (not `model.generate()`) so the penalty is applied exactly as specified, without HuggingFace's internal logic.
- Works with any model that `AutoModelForCausalLM` can load.