File size: 2,361 Bytes
6e802a4 | 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 74 75 76 77 78 79 | # test_stop_supervision.py: The stop-head supervision.
import pathlib
import sys
import pytest
import torch
sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent))
IGNORE_INDEX = -100
def _tok():
from tokenizer import TokenRegistry, expand_tokenizer, load_base_tokenizer
tokenizer = load_base_tokenizer()
expand_tokenizer(tokenizer)
return tokenizer, TokenRegistry(tokenizer)
@pytest.fixture(scope="module")
def build():
from delay_dataset import DelayMimiDataset
def _make(tmpdir, weight):
import json
tok, reg = _tok()
p = pathlib.Path(tmpdir) / "rows.jsonl"
nlev = __import__("delay_dataset").NUM_LEVELS
row = {"prompt": "hello there friend",
"speech_tokens": [5, 6, 7, 8],
"residual_codes": [[1, 2, 3, 4] for _ in range(nlev - 1)],
"speaker_id": 1}
p.write_text(json.dumps(row) + "\n", encoding="utf-8")
return DelayMimiDataset(str(p), tok, reg, stage=1, lazy=False,
audio_continue_weight=weight)
return _make
def _audio_span(item, reg):
ids = item["input_ids"].tolist()
start = ids.index(reg.audio_start_id) + 1
end = ids.index(reg.audio_end_id)
return start, end
def test_default_is_the_old_behaviour(build, tmp_path):
_, reg = _tok()
it = build(tmp_path, 0.0)[0]
s, e = _audio_span(it, reg)
assert (it["labels"][s:e] == IGNORE_INDEX).all()
def test_enabled_supervises_every_audio_frame(build, tmp_path):
_, reg = _tok()
it = build(tmp_path, 0.1)[0]
s, e = _audio_span(it, reg)
assert e > s
assert (it["labels"][s:e] != IGNORE_INDEX).all()
assert (it["labels"][s:e] == reg.audio_start_id).all()
def test_stop_token_stays_supervised_either_way(build, tmp_path):
_, reg = _tok()
for w in (0.0, 0.1):
it = build(tmp_path, w)[0]
ids = it["input_ids"].tolist()
assert it["labels"][ids.index(reg.audio_end_id)] == reg.audio_end_id
def test_audio_weight_is_small_relative_to_text(build, tmp_path):
_, reg = _tok()
it = build(tmp_path, 0.1)[0]
s, e = _audio_span(it, reg)
lw = it["loss_weights"]
assert torch.allclose(lw[s:e], torch.tensor(0.1))
assert lw[ids_end := it["input_ids"].tolist().index(reg.audio_end_id)] == 1.0
assert lw[ids_end] > lw[s]
|