| |
|
|
| 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] |
|
|
