import unittest import torch import numpy as np from pathlib import Path import tempfile import sys import os sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) from utils.bct import Block, BlockSequence, BlockType from block_types import CondAudioTextBlockType, CondAudioBlockType, CausalSemanticBlockType class TestBCTTextBlocks(unittest.TestCase): """Test Block Conditioning Transforms (BCT) functionality with text blocks.""" def setUp(self): """Set up test fixtures.""" torch.manual_seed(42) np.random.seed(42) # Create sample text description block self.text_block = Block( spec=CondAudioTextBlockType, inputs={"text_input": torch.tensor([1, 2, 3, 4])}, debug_text="[extract drums]", ) # Create sample conditioning block using CondAudioBlockType # CondAudioTextBlockType should be paired with CondAudioBlockType for unified conditioning self.conditioning_block = Block( spec=CondAudioBlockType, inputs={"semantic_input": torch.randint(0, 4096, (50, 1)).to(torch.float16)}, ) # Create sample semantic block self.semantic_block = Block( spec=CausalSemanticBlockType, inputs={"semantic_input": torch.randint(0, 4096, (100, 1)).to(torch.float16)}, targets={"semantic_output": torch.randint(0, 4096, (100, 1)).to(torch.float16)}, ) def test_text_description_block_creation(self): """Test creation and properties of text description blocks.""" self.assertEqual(self.text_block.spec, CondAudioTextBlockType) self.assertEqual(self.text_block.spec.name, "text_description") self.assertTrue(self.text_block.spec.is_causal) self.assertEqual(len(self.text_block), 4) # Length should be from text_input self.assertEqual(self.text_block.debug_text, "[extract drums]") def test_text_conditioning_pair_in_block_sequence(self): """Test text/conditioning pairs in block sequences.""" # Create a block sequence with text/conditioning pair blocks = [self.text_block, self.conditioning_block, self.semantic_block] sequence = BlockSequence(blocks) self.assertEqual(len(sequence), 3) self.assertEqual(sequence.n_tokens, 4 + 50 + 100) # Sum of all block lengths # Verify order is maintained self.assertEqual(sequence[0].spec.name, "text_description") self.assertEqual(sequence[1].spec.name, "cond_audio") self.assertEqual(sequence[2].spec.name, "semantic") def test_block_sequence_string_representation(self): """Test string representation includes debug text.""" sequence = BlockSequence([self.text_block, self.conditioning_block]) # Print the sequence for visual verification print("\n=== Block Sequence String Representation ===") print(sequence) print("============================================\n") sequence_str = str(sequence) self.assertIn("text_description", sequence_str) self.assertIn("cond_audio", sequence_str) self.assertIn("[extract drums]", sequence_str) # Debug text should appear def test_multiple_text_conditioning_pairs(self): """Test sequence with multiple text/conditioning pairs.""" # Create additional text blocks with different content artist_text_block = Block( spec=CondAudioTextBlockType, inputs={"text_input": torch.tensor([5, 6, 7])}, debug_text="[in the style of this artist]", ) artist_conditioning_block = Block( spec=CausalSemanticBlockType, # Using different block type inputs={"semantic_input": torch.randint(0, 4096, (30, 1))}, ) # Create sequence: text1, conditioning1, text2, conditioning2, main_content blocks = [ self.text_block, # [extract drums] self.conditioning_block, # stem content artist_text_block, # [in the style of this artist] artist_conditioning_block, # artist content self.semantic_block, # main generation target ] sequence = BlockSequence(blocks) print("\n=== Multiple Text/Conditioning Pairs ===") print(f"Total blocks: {len(sequence)}") print(f"Total tokens: {sequence.n_tokens}") for i, block in enumerate(sequence): if hasattr(block, "debug_text") and block.debug_text: print(f"Block {i}: {block.spec.name} ({len(block)} tokens) - '{block.debug_text}'") else: print(f"Block {i}: {block.spec.name} ({len(block)} tokens)") print("=======================================\n") # Verify structure self.assertEqual(len(sequence), 5) self.assertEqual(sequence[0].spec.name, "text_description") # Changed to use unified cond_audio block type for conditioning self.assertEqual(sequence[1].spec.name, "cond_audio") self.assertEqual(sequence[2].spec.name, "text_description") self.assertEqual(sequence[3].spec.name, "semantic") # artist conditioning self.assertEqual(sequence[4].spec.name, "semantic") # main content def test_block_sequence_cropping_with_text_blocks(self): """Test that cropping works correctly with text blocks.""" # Create sequence with known token counts sequence = BlockSequence( [ self.text_block, # 4 tokens self.conditioning_block, # 50 tokens self.semantic_block, # 100 tokens ] ) original_total = sequence.n_tokens # 154 tokens # Crop to 60 tokens (should include text + conditioning + partial semantic) cropped = sequence.crop_to_max_tokens(60) self.assertEqual(cropped.n_tokens, 60) self.assertEqual(len(cropped), 3) # All blocks should be present # First two blocks should be unchanged self.assertEqual(len(cropped[0]), 4) # text block unchanged self.assertEqual(len(cropped[1]), 50) # conditioning block unchanged self.assertEqual(len(cropped[2]), 6) # semantic block truncated to 60-4-50=6 def test_block_sequence_serialization_with_text_blocks(self): """Test saving and loading sequences with text blocks.""" sequence = BlockSequence([self.text_block, self.conditioning_block]) with tempfile.TemporaryDirectory() as temp_dir: save_path = Path(temp_dir) / "test_sequence.pt" # Save sequence sequence.save(save_path) self.assertTrue(save_path.exists()) # Load sequence loaded_sequence = BlockSequence.load(save_path) # Verify structure is preserved self.assertEqual(len(loaded_sequence), len(sequence)) # Verify text block is preserved loaded_text_block = loaded_sequence[0] self.assertEqual(loaded_text_block.spec.name, "text_description") self.assertTrue( torch.equal(loaded_text_block.inputs["text_input"], self.text_block.inputs["text_input"]) ) # Note: debug_text is not preserved in serialization (by design) # Verify conditioning block is preserved loaded_conditioning_block = loaded_sequence[1] self.assertEqual(loaded_conditioning_block.spec.name, "cond_audio") def test_block_pairing_pattern_verification(self): """Test the specific pairing pattern used in data_utils.py.""" # Simulate the exact pattern from make_text_conditioning_pair pairs = [] # Create multiple conditioning types as tuples conditioning_types = [ ("prefix", "[continue from this]"), ("stem", "[extract Bass, Drums]"), ("suffix", "[lead up to this]"), ] for block_type, debug_text in conditioning_types: text_block = Block( spec=CondAudioTextBlockType, inputs={"text_input": torch.tensor([1, 2, 3])}, debug_text=debug_text, ) conditioning_block = Block( # Use unified CondAudioBlockType for all text conditioning pairs spec=CondAudioBlockType, inputs={"semantic_input": torch.randint(0, 4096, (25, 1))}, ) pairs.append((text_block, conditioning_block)) # Flatten pairs like in data_utils.py flattened_blocks = [] for block in pairs: if isinstance(block, tuple): text_block, content_block = block flattened_blocks.extend([text_block, content_block]) else: flattened_blocks.append(block) sequence = BlockSequence(flattened_blocks) print("\n=== Pairing Pattern Verification ===") print(f"Original pairs: {len(pairs)}") print(f"Flattened blocks: {len(sequence)}") for i, block in enumerate(sequence): pair_num = i // 2 position = "text" if i % 2 == 0 else "content" debug_info = ( f" - '{block.debug_text}'" if hasattr(block, "debug_text") and block.debug_text else "" ) print(f"Block {i}: Pair {pair_num} {position} ({block.spec.name}){debug_info}") print("====================================\n") # Verify alternating pattern: text, content, text, content, text, content self.assertEqual(len(sequence), 6) for i in range(0, len(sequence), 2): self.assertEqual(sequence[i].spec.name, "text_description") if i + 1 < len(sequence): self.assertNotEqual(sequence[i + 1].spec.name, "text_description") def test_empty_and_edge_cases(self): """Test edge cases with text blocks.""" # Empty sequence empty_sequence = BlockSequence([]) self.assertEqual(len(empty_sequence), 0) self.assertEqual(empty_sequence.n_tokens, 0) # Single text block single_sequence = BlockSequence([self.text_block]) self.assertEqual(len(single_sequence), 1) self.assertEqual(single_sequence.n_tokens, 4) # Text block with empty inputs empty_text_block = Block( spec=CondAudioTextBlockType, inputs={"text_input": torch.tensor([])}, # Empty tensor debug_text="[empty]", ) self.assertEqual(len(empty_text_block), 0) if __name__ == "__main__": unittest.main()