| import json |
| import unittest |
| from pathlib import Path |
| from tempfile import TemporaryDirectory |
|
|
| import torch |
|
|
| from speech_bridge_gemma.slice_longcat_codes import slice_codes |
|
|
|
|
| class SliceLongCatCodesTest(unittest.TestCase): |
| def test_slice_uses_manifest_subset_index(self) -> None: |
| with TemporaryDirectory() as tmp: |
| root = Path(tmp) |
| source = root / "source.pt" |
| manifest = root / "manifest.jsonl" |
| out = root / "out.pt" |
| torch.save( |
| { |
| "meta": {"codec": "longcat", "n_acoustic_codebooks": 1, "input_sample_rate": 16000}, |
| "semantic_codes": [ |
| torch.tensor([0, 0]), |
| torch.tensor([1, 1, 1]), |
| torch.tensor([2]), |
| ], |
| "acoustic_codes": [ |
| torch.tensor([[10, 10]]), |
| torch.tensor([[11, 11, 11]]), |
| torch.tensor([[12]]), |
| ], |
| }, |
| source, |
| ) |
| rows = [ |
| {"id": "row_b", "text": "texto b", "prompt": "pergunta b", "answer": "resposta b", "subset_index": 2, "duration": 1.0}, |
| {"id": "row_a", "text": "texto a", "question": "pergunta a", "subset_index": 0, "duration_sec": 2.0}, |
| ] |
| manifest.write_text("\n".join(json.dumps(row) for row in rows) + "\n", encoding="utf-8") |
|
|
| summary = slice_codes(source, manifest, out, "subset_index") |
| payload = torch.load(out, map_location="cpu", weights_only=False) |
|
|
| self.assertEqual(summary["rows"], 2) |
| self.assertEqual(payload["semantic_codes"][0].tolist(), [2]) |
| self.assertEqual(payload["semantic_codes"][1].tolist(), [0, 0]) |
| self.assertEqual(payload["acoustic_codes"][0].tolist(), [[12]]) |
| self.assertEqual(payload["meta"]["samples"][0]["id"], "row_b") |
| self.assertEqual(payload["meta"]["samples"][0]["frames"], 1) |
| self.assertEqual(payload["meta"]["samples"][1]["answer_text"], "texto a") |
|
|
|
|
| if __name__ == "__main__": |
| unittest.main() |
|
|
