# test_dataset.py: JSONL parsing, sequence construction and collation. import json import os import tempfile import pytest import torch from config import ( AUDIO_END_TOKEN, AUDIO_START_TOKEN, NUM_SPEECH_TOKENS, THINK_END_TOKEN, THINK_START_TOKEN, USER_PROMPT_TOKEN, get_all_special_tokens, ) from dataset import MambaCoTDataCollator class SimpleTokenizer: def __init__(self): self.vocab = {} self._next_id = 0 self.eos_token_id = self._add("[EOS]") for token in get_all_special_tokens(): self._add(token) def _add(self, token): if token not in self.vocab: self.vocab[token] = self._next_id self._next_id += 1 return self.vocab[token] def encode(self, text, add_special_tokens=False, return_tensors=None): tokens = text.split() ids = [self.vocab.get(t, self._add(t)) for t in tokens] if return_tensors == "pt": return torch.tensor([ids], dtype=torch.long) return ids def convert_tokens_to_ids(self, tokens): return [self.vocab.get(t, self._add(t)) for t in tokens] def convert_ids_to_tokens(self, ids): inv = {v: k for k, v in self.vocab.items()} return [inv.get(i, f"[UNK_{i}]") for i in ids] def decode(self, ids, skip_special_tokens=False): return " ".join(self.convert_ids_to_tokens(ids)) def __len__(self): return len(self.vocab) def add_special_tokens(self, special_tokens_dict): count = 0 for tokens in special_tokens_dict.values(): if isinstance(tokens, list): for t in tokens: if t not in self.vocab: self._add(t) count += 1 return count class TestJSONLParsing: def test_valid_jsonl(self): data = [ { "prompt": "Hello world", "reasoning": "Neutral greeting", "audio_file": "test.wav", }, { "prompt": "Oh great", "reasoning": "Sarcastic [EMO:sarcastic]", "audio_file": "test2.wav", "emotion_label": "sarcastic", }, ] with tempfile.NamedTemporaryFile(mode="w", suffix=".jsonl", delete=False) as f: for item in data: f.write(json.dumps(item) + "\n") tmpfile = f.name try: parsed = [] with open(tmpfile, "r") as f: for line in f: if line.strip(): parsed.append(json.loads(line)) assert len(parsed) == 2 assert parsed[0]["prompt"] == "Hello world" assert parsed[1]["emotion_label"] == "sarcastic" finally: os.unlink(tmpfile) def test_empty_lines_skipped(self): with tempfile.NamedTemporaryFile(mode="w", suffix=".jsonl", delete=False) as f: f.write(json.dumps({"prompt": "a", "reasoning": "", "audio_file": "x.wav"}) + "\n") f.write("\n") f.write(" \n") f.write(json.dumps({"prompt": "b", "reasoning": "", "audio_file": "y.wav"}) + "\n") tmpfile = f.name try: parsed = [] with open(tmpfile, "r") as f: for line in f: if line.strip(): parsed.append(json.loads(line)) assert len(parsed) == 2 finally: os.unlink(tmpfile) def test_optional_fields(self): data = {"prompt": "Test", "reasoning": "", "audio_file": "x.wav", "speech_tokens": [42, 100, 200]} with tempfile.NamedTemporaryFile(mode="w", suffix=".jsonl", delete=False) as f: f.write(json.dumps(data) + "\n") tmpfile = f.name try: with open(tmpfile, "r") as f: item = json.loads(f.readline()) assert item["speech_tokens"] == [42, 100, 200] finally: os.unlink(tmpfile) class TestSequenceConstruction: def setup_method(self): self.tokenizer = SimpleTokenizer() def test_stage1_no_reasoning(self): prompt = "Hello world" prompt_text = f"{USER_PROMPT_TOKEN} {prompt}" prompt_ids = self.tokenizer.encode(prompt_text) audio_start_ids = self.tokenizer.encode(AUDIO_START_TOKEN) speech_ids = self.tokenizer.convert_tokens_to_ids(["[SPEECH_42]", "[SPEECH_100]"]) audio_end_ids = self.tokenizer.encode(AUDIO_END_TOKEN) input_ids = prompt_ids + audio_start_ids + speech_ids + audio_end_ids input_ids.append(self.tokenizer.eos_token_id) tokens = self.tokenizer.convert_ids_to_tokens(input_ids) assert USER_PROMPT_TOKEN in tokens assert AUDIO_START_TOKEN in tokens assert AUDIO_END_TOKEN in tokens assert THINK_START_TOKEN not in tokens assert THINK_END_TOKEN not in tokens def test_stage2_with_reasoning(self): prompt = "Great Monday" reasoning = "[EMO:sarcastic] [PACE:slow]" prompt_ids = self.tokenizer.encode(f"{USER_PROMPT_TOKEN} {prompt}") reasoning_ids = self.tokenizer.encode(f"{THINK_START_TOKEN} {reasoning} {THINK_END_TOKEN}") audio_start_ids = self.tokenizer.encode(AUDIO_START_TOKEN) speech_ids = self.tokenizer.convert_tokens_to_ids(["[SPEECH_0]"]) audio_end_ids = self.tokenizer.encode(AUDIO_END_TOKEN) input_ids = prompt_ids + reasoning_ids + audio_start_ids + speech_ids + audio_end_ids input_ids.append(self.tokenizer.eos_token_id) tokens = self.tokenizer.convert_ids_to_tokens(input_ids) assert THINK_START_TOKEN in tokens assert THINK_END_TOKEN in tokens assert "[EMO:sarcastic]" in tokens def test_speech_token_range(self): speech_tokens = [f"[SPEECH_{i}]" for i in [0, 512, NUM_SPEECH_TOKENS - 1]] ids = self.tokenizer.convert_tokens_to_ids(speech_tokens) assert len(ids) == 3 assert all(isinstance(i, int) for i in ids) @pytest.fixture(scope="module") def real_registry(): from tokenizer import load_base_tokenizer, expand_tokenizer, TokenRegistry tok = load_base_tokenizer() expand_tokenizer(tok) return TokenRegistry(tok) class TestCollator: def _make_item(self, registry, input_ids, prompt_end, reasoning_start=0, reasoning_end=0, speech_start=0, speech_end=0, residual_codes=None, speaker_id=None): return { "input_ids": torch.tensor(input_ids, dtype=torch.long), "prompt_end": prompt_end, "reasoning_start": reasoning_start, "reasoning_end": reasoning_end, "speech_start": speech_start, "speech_end": speech_end, "residual_codes": residual_codes, "speaker_id": speaker_id, } def test_padding(self, real_registry): collator = MambaCoTDataCollator( pad_token_id=real_registry.pad_id, registry=real_registry, stage=1, ) batch = [ self._make_item( real_registry, input_ids=[real_registry.bos_id, 100, 101, real_registry.eos_id], prompt_end=2, speech_start=2, speech_end=4, ), self._make_item( real_registry, input_ids=[real_registry.bos_id, 100, 101, 102, 103, real_registry.eos_id], prompt_end=2, speech_start=2, speech_end=6, ), ] result = collator(batch) assert result["input_ids"].shape == (2, 6) assert result["input_ids"][0, 4].item() == real_registry.pad_id assert result["input_ids"][0, 5].item() == real_registry.pad_id def test_prompt_masked_in_labels(self, real_registry): collator = MambaCoTDataCollator( pad_token_id=real_registry.pad_id, registry=real_registry, stage=1, ) batch = [ self._make_item( real_registry, input_ids=[real_registry.bos_id, 100, 101, 102, real_registry.eos_id], prompt_end=3, speech_start=3, speech_end=5, ), ] result = collator(batch) labels = result["labels"] assert labels[0, 0].item() == -100 assert labels[0, 1].item() == -100 assert labels[0, 2].item() == -100 assert labels[0, 3].item() != -100 def test_padding_masked_in_labels(self, real_registry): collator = MambaCoTDataCollator( pad_token_id=real_registry.pad_id, registry=real_registry, stage=1, ) batch = [ self._make_item( real_registry, input_ids=[real_registry.bos_id, 100, real_registry.eos_id], prompt_end=1, speech_start=1, speech_end=3, ), self._make_item( real_registry, input_ids=[real_registry.bos_id, 100, 101, 102, real_registry.eos_id], prompt_end=1, speech_start=1, speech_end=5, ), ] result = collator(batch) labels = result["labels"] assert labels[0, 3].item() == -100 assert labels[0, 4].item() == -100 def test_attention_mask(self, real_registry): collator = MambaCoTDataCollator( pad_token_id=real_registry.pad_id, registry=real_registry, stage=1, ) batch = [ self._make_item( real_registry, input_ids=[real_registry.bos_id, 100, real_registry.eos_id], prompt_end=1, speech_start=1, speech_end=3, ), self._make_item( real_registry, input_ids=[real_registry.bos_id, 100, 101, 102, real_registry.eos_id], prompt_end=1, speech_start=1, speech_end=5, ), ] result = collator(batch) mask = result["attention_mask"] assert mask[0, 0].item() == 1 assert mask[0, 2].item() == 1 assert mask[0, 3].item() == 0 assert mask[0, 4].item() == 0 assert mask[1].sum().item() == 5 def test_loss_weights_stage1(self, real_registry): collator = MambaCoTDataCollator( pad_token_id=real_registry.pad_id, registry=real_registry, stage=1, ) batch = [ self._make_item( real_registry, input_ids=[real_registry.bos_id, 100, 101, 102, real_registry.eos_id], prompt_end=2, reasoning_start=2, reasoning_end=2, speech_start=2, speech_end=5, ), ] result = collator(batch) weights = result["loss_weights"] assert weights[0, 0].item() == pytest.approx(0.0) assert weights[0, 1].item() == pytest.approx(0.0) assert weights[0, 2].item() == pytest.approx(1.0) assert weights[0, 3].item() == pytest.approx(1.0) def test_loss_weights_stage2(self, real_registry): lambda_r = 0.3 lambda_s = 1.0 collator = MambaCoTDataCollator( pad_token_id=real_registry.pad_id, registry=real_registry, stage=2, lambda_r=lambda_r, lambda_s=lambda_s, ) batch = [ self._make_item( real_registry, input_ids=[ real_registry.bos_id, real_registry.user_prompt_id, 100, real_registry.think_start_id, 200, 201, real_registry.think_end_id, real_registry.audio_start_id, 300, real_registry.audio_end_id, real_registry.eos_id, ], prompt_end=3, reasoning_start=4, reasoning_end=6, speech_start=8, speech_end=9, ), ] result = collator(batch) weights = result["loss_weights"] assert weights[0, 0].item() == pytest.approx(0.0) assert weights[0, 1].item() == pytest.approx(0.0) assert weights[0, 2].item() == pytest.approx(0.0) assert weights[0, 3].item() == pytest.approx(lambda_r) assert weights[0, 4].item() == pytest.approx(lambda_r) assert weights[0, 5].item() == pytest.approx(lambda_r) assert weights[0, 7].item() == pytest.approx(lambda_s) assert weights[0, 8].item() == pytest.approx(lambda_s) def test_residual_codes_none(self, real_registry): collator = MambaCoTDataCollator( pad_token_id=real_registry.pad_id, registry=real_registry, stage=1, ) batch = [ self._make_item( real_registry, input_ids=[real_registry.bos_id, 100, real_registry.eos_id], prompt_end=1, speech_start=1, speech_end=3, residual_codes=None, ), ] result = collator(batch) assert result["residual_codes"] is None def test_speaker_ids_none(self, real_registry): collator = MambaCoTDataCollator( pad_token_id=real_registry.pad_id, registry=real_registry, stage=1, ) batch = [ self._make_item( real_registry, input_ids=[real_registry.bos_id, 100, real_registry.eos_id], prompt_end=1, speech_start=1, speech_end=3, speaker_id=None, ), ] result = collator(batch) assert result["speaker_ids"] is None