import unittest import random import numpy as np import torch from unittest.mock import Mock, patch import sys import os sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) from data_utils import make_text_conditioning_pair, INSTRUCTION_TEMPLATES from block_types import CondAudioTextBlockType from utils.bct import Block, BlockSequence class TestTextConditioningPairs(unittest.TestCase): def setUp(self): """Set up test fixtures.""" random.seed(42) np.random.seed(42) torch.manual_seed(42) # Mock tokenizer path self.tokenizer_fp = "/path/to/tokenizer" # Create test data tensors self.semantic_data = torch.randint(0, 4096, (1, 100)) # Semantic data self.vae_data = torch.randn(100, 128) # VAE data @patch("data_utils.tokenize_batch") @patch("data_utils.make_conditioning_block") def test_make_text_conditioning_pair_basic(self, mock_make_conditioning, mock_tokenize): """Test basic text conditioning pair creation.""" # Mock tokenization mock_tokenize.return_value = torch.tensor([1, 2, 3, 4]) # Mock conditioning block mock_conditioning_block = Mock() mock_make_conditioning.return_value = mock_conditioning_block # Test with basic block type text_block, content_block = make_text_conditioning_pair( data=self.semantic_data, block_type="artist", is_diffusion=False, tokenizer_fp=self.tokenizer_fp, ) # Verify text block structure self.assertEqual(text_block.spec, CondAudioTextBlockType) self.assertTrue("text_input" in text_block.inputs) self.assertIsNotNone(text_block.debug_text) self.assertTrue(text_block.debug_text.startswith("[") and text_block.debug_text.endswith("]")) # Verify content block is from make_conditioning_block self.assertEqual(content_block, mock_conditioning_block) # Verify function calls mock_tokenize.assert_called_once() mock_make_conditioning.assert_called_once_with( self.semantic_data, "cond_audio", False, False, False, debug_text="artist" ) @patch("data_utils.tokenize_batch") @patch("data_utils.make_conditioning_block") def test_stem_type_integration(self, mock_make_conditioning, mock_tokenize): """Test stem_type integration in text descriptions.""" mock_tokenize.return_value = torch.tensor([1, 2, 3, 4, 5]) mock_make_conditioning.return_value = Mock() # Test with stem_type text_block, content_block = make_text_conditioning_pair( data=self.semantic_data, block_type="stem", is_diffusion=False, tokenizer_fp=self.tokenizer_fp, stem_type="add Bass, Drums", ) # Verify stem_type is included in text description debug_text = text_block.debug_text self.assertIn("Bass, Drums", debug_text) self.assertTrue(debug_text.startswith("[") and debug_text.endswith("]")) @patch("data_utils.tokenize_batch") @patch("data_utils.make_conditioning_block") def test_instruction_template_usage(self, mock_make_conditioning, mock_tokenize): """Test that instruction templates are used correctly.""" mock_tokenize.return_value = torch.tensor([1, 2, 3]) mock_make_conditioning.return_value = Mock() # Test multiple calls to see template variation debug_texts = [] for _ in range(10): text_block, _ = make_text_conditioning_pair( data=self.semantic_data, block_type="cover", is_diffusion=False, tokenizer_fp=self.tokenizer_fp, ) debug_texts.append(text_block.debug_text) # Should use templates from INSTRUCTION_TEMPLATES["cover"] for text in debug_texts: self.assertTrue(text.startswith("[") and text.endswith("]")) text_content = text[1:-1] # Remove brackets # Should be one of the cover templates self.assertTrue(any(template == text_content for template in INSTRUCTION_TEMPLATES["cover"])) @patch("data_utils.tokenize_batch") @patch("data_utils.make_conditioning_block") def test_diffusion_vs_semantic(self, mock_make_conditioning, mock_tokenize): """Test handling of diffusion vs semantic data.""" mock_tokenize.return_value = torch.tensor([1, 2, 3]) mock_make_conditioning.return_value = Mock() # Test semantic make_text_conditioning_pair( data=self.semantic_data, block_type="vox", is_diffusion=False, tokenizer_fp=self.tokenizer_fp ) mock_make_conditioning.assert_called_with( self.semantic_data, "cond_audio", False, False, False, debug_text="vox" ) # Test diffusion make_text_conditioning_pair( data=self.vae_data, block_type="vox", is_diffusion=True, tokenizer_fp=self.tokenizer_fp ) mock_make_conditioning.assert_called_with( self.vae_data, "cond_audio", True, False, False, debug_text="vox" ) def test_instruction_templates_completeness(self): """Test that all expected block types have instruction templates.""" expected_block_types = [ "stem_add", "stem_extract", "stem_remove", "cover", "artist", "playlist", "vox", "overpaint", "underpaint", "sample", "prefix", "suffix", ] for block_type in expected_block_types: self.assertIn(block_type, INSTRUCTION_TEMPLATES) self.assertIsInstance(INSTRUCTION_TEMPLATES[block_type], list) self.assertGreater(len(INSTRUCTION_TEMPLATES[block_type]), 0) # Each template should be a string for template in INSTRUCTION_TEMPLATES[block_type]: self.assertIsInstance(template, str) self.assertGreater(len(template.strip()), 0) def test_block_sequence_pairing_order(self): """Test that text/conditioning pairs maintain correct order in block sequences.""" # Create test data for multiple conditioning types test_pairs = [] with ( patch("data_utils.tokenize_batch") as mock_tokenize, patch("data_utils.make_conditioning_block") as mock_conditioning, ): mock_tokenize.return_value = torch.tensor([1, 2, 3]) # Create different types of conditioning pairs conditioning_types = ["artist", "stem", "vox", "prefix"] for i, ctype in enumerate(conditioning_types): # Mock different conditioning blocks for each type mock_conditioning_block = Mock() mock_conditioning_block.spec = Mock() mock_conditioning_block.spec.name = f"mock_{ctype}" mock_conditioning.return_value = mock_conditioning_block pair = make_text_conditioning_pair( data=self.semantic_data, block_type=ctype, is_diffusion=False, tokenizer_fp=self.tokenizer_fp, stem_type="add Bass" if ctype == "stem" else None, ) test_pairs.append(pair) # Simulate the flattening logic from data_utils.py flattened_blocks = [] for block in test_pairs: if isinstance(block, tuple): text_block, content_block = block flattened_blocks.extend([text_block, content_block]) else: flattened_blocks.append(block) # Create block sequence block_sequence = BlockSequence(flattened_blocks) # Print block sequence for visual verification print("\n=== Block Sequence Structure ===") print(f"Total blocks: {len(block_sequence.blocks)}") for i, block in enumerate(block_sequence.blocks): if hasattr(block, "debug_text"): print(f"Block {i}: {block.spec.name} - '{block.debug_text}'") else: print(f"Block {i}: {block.spec.name}") print("================================\n") # Verify pairing order: text_description, then conditioning content self.assertEqual(len(block_sequence.blocks), 8) # 4 pairs = 8 blocks for i in range(0, len(block_sequence.blocks), 2): # Even indices should be text description blocks text_block = block_sequence.blocks[i] self.assertEqual(text_block.spec.name, "text_description") self.assertIsNotNone(text_block.debug_text) # Odd indices should be conditioning content blocks if i + 1 < len(block_sequence.blocks): content_block = block_sequence.blocks[i + 1] self.assertNotEqual(content_block.spec.name, "text_description") @patch("data_utils.tokenize_batch") @patch("data_utils.make_conditioning_block") def test_fallback_behavior(self, mock_make_conditioning, mock_tokenize): """Test fallback behavior for unknown block types.""" mock_tokenize.return_value = torch.tensor([1, 2]) mock_make_conditioning.return_value = Mock() # Test with unknown block type text_block, _ = make_text_conditioning_pair( data=self.semantic_data, block_type="unknown_type", is_diffusion=False, tokenizer_fp=self.tokenizer_fp, ) # Should fallback to simple block type text self.assertEqual(text_block.debug_text, "[unknown_type]") def test_stem_type_format_handling(self): """Test various stem_type formats.""" with ( patch("data_utils.tokenize_batch") as mock_tokenize, patch("data_utils.make_conditioning_block") as mock_conditioning, ): mock_tokenize.return_value = torch.tensor([1, 2, 3]) mock_conditioning.return_value = Mock() # Test proper "add " format text_block, _ = make_text_conditioning_pair( data=self.semantic_data, block_type="stem", is_diffusion=False, tokenizer_fp=self.tokenizer_fp, stem_type="add Bass, Drums", ) self.assertIn("Bass, Drums", text_block.debug_text) # Test improper format (should use original) text_block, _ = make_text_conditioning_pair( data=self.semantic_data, block_type="stem", is_diffusion=False, tokenizer_fp=self.tokenizer_fp, stem_type="extract vocals", ) self.assertEqual(text_block.debug_text, "[extract vocals]") if __name__ == "__main__": unittest.main()