168 lines
4.7 KiB
Python
168 lines
4.7 KiB
Python
import re
|
|
from unittest.mock import MagicMock
|
|
|
|
import pytest
|
|
import torch
|
|
from src.nodes.break_encoder import BreakEncoder
|
|
from src.utils import clean_prompt
|
|
|
|
|
|
def split_blocks(text, keyword, cleanup=True):
|
|
"""Mirror of the encode() split logic."""
|
|
blocks = re.split(rf"\s*{re.escape(keyword)}\s*", text)
|
|
blocks = [b.strip() for b in blocks]
|
|
blocks = [b for b in blocks if b]
|
|
if cleanup:
|
|
blocks = [clean_prompt(b) for b in blocks]
|
|
blocks = [b for b in blocks if b.strip()]
|
|
return blocks
|
|
|
|
|
|
@pytest.fixture
|
|
def node():
|
|
return BreakEncoder()
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_clip():
|
|
cond = torch.zeros(1, 10, 768)
|
|
pooled = torch.zeros(1, 768)
|
|
clip = MagicMock()
|
|
clip.tokenize.return_value = MagicMock()
|
|
clip.encode_from_tokens_scheduled.return_value = [[cond, {"pooled_output": pooled}]]
|
|
return clip
|
|
|
|
|
|
def run(node, clip, text, keyword="BREAK", cleanup=False):
|
|
return node.encode(clip, keyword, text, cleanup)
|
|
|
|
|
|
# --- Split logic ---
|
|
|
|
|
|
def test_split_basic():
|
|
blocks = split_blocks("tag1\nBREAK\ntag2", "BREAK", cleanup=False)
|
|
assert blocks == ["tag1", "tag2"]
|
|
|
|
|
|
def test_split_multiple_breaks():
|
|
blocks = split_blocks("a\nBREAK\nb\nBREAK\nc", "BREAK", cleanup=False)
|
|
assert blocks == ["a", "b", "c"]
|
|
|
|
|
|
def test_split_break_inline():
|
|
blocks = split_blocks("tag1, BREAK, tag2", "BREAK", cleanup=False)
|
|
assert blocks == ["tag1,", ", tag2"]
|
|
|
|
|
|
def test_split_custom_keyword():
|
|
blocks = split_blocks("a\n---\nb", "---", cleanup=False)
|
|
assert blocks == ["a", "b"]
|
|
|
|
|
|
def test_split_case_sensitive():
|
|
# "break" lowercase should NOT match "BREAK"
|
|
blocks = split_blocks("a\nbreak\nb", "BREAK", cleanup=False)
|
|
assert blocks == ["a\nbreak\nb"]
|
|
|
|
|
|
def test_split_filters_empty_blocks():
|
|
blocks = split_blocks("a\nBREAK\n\nBREAK\nb", "BREAK", cleanup=False)
|
|
assert blocks == ["a", "b"]
|
|
|
|
|
|
def test_split_break_at_start():
|
|
blocks = split_blocks("BREAK\na\nBREAK\nb", "BREAK", cleanup=False)
|
|
assert blocks == ["a", "b"]
|
|
|
|
|
|
def test_split_break_at_end():
|
|
blocks = split_blocks("a\nBREAK\nb\nBREAK", "BREAK", cleanup=False)
|
|
assert blocks == ["a", "b"]
|
|
|
|
|
|
def test_split_whitespace_around_break():
|
|
blocks = split_blocks("a \n BREAK \n b", "BREAK", cleanup=False)
|
|
assert blocks == ["a", "b"]
|
|
|
|
|
|
# --- Cleanup integration ---
|
|
|
|
|
|
def test_split_cleanup_removes_double_commas():
|
|
blocks = split_blocks("a,,b\nBREAK\nc", "BREAK", cleanup=True)
|
|
assert ",," not in blocks[0]
|
|
|
|
|
|
def test_split_cleanup_empty_after_clean():
|
|
# A block that becomes empty after cleanup should be filtered
|
|
blocks = split_blocks("valid\nBREAK\n,,\nBREAK\nalso valid", "BREAK", cleanup=True)
|
|
assert all(b.strip() for b in blocks)
|
|
|
|
|
|
# --- Validation ---
|
|
|
|
|
|
def test_empty_text_raises(node, mock_clip):
|
|
with pytest.raises(ValueError, match="empty"):
|
|
run(node, mock_clip, " ")
|
|
|
|
|
|
def test_none_clip_raises(node):
|
|
with pytest.raises(ValueError, match="CLIP"):
|
|
run(node, None, "some text")
|
|
|
|
|
|
def test_empty_keyword_raises(node, mock_clip):
|
|
with pytest.raises(ValueError, match="keyword"):
|
|
run(node, mock_clip, "some text", keyword=" ")
|
|
|
|
|
|
def test_all_blocks_empty_after_cleanup_raises(node, mock_clip):
|
|
with pytest.raises(ValueError, match="empty after cleanup"):
|
|
run(node, mock_clip, ",, BREAK ,,", cleanup=True)
|
|
|
|
|
|
# --- Encoding ---
|
|
|
|
|
|
def test_encode_single_block(node, mock_clip):
|
|
conditioning, text = run(node, mock_clip, "character, solo")
|
|
assert conditioning is not None
|
|
assert text == "character, solo"
|
|
|
|
|
|
def test_encode_two_blocks_concat(node):
|
|
cond1 = MagicMock()
|
|
cond1.shape = (1, 10, 768)
|
|
cond2 = MagicMock()
|
|
cond2.shape = (1, 5, 768)
|
|
pooled = MagicMock()
|
|
|
|
# Configure torch.cat
|
|
concat_result = MagicMock()
|
|
concat_result.shape = [1, 15, 768]
|
|
torch.cat.return_value = concat_result
|
|
|
|
clip = MagicMock()
|
|
clip.tokenize.return_value = MagicMock()
|
|
clip.encode_from_tokens_scheduled.side_effect = [
|
|
[[cond1, {"pooled_output": pooled}]],
|
|
[[cond2, {"pooled_output": pooled}]],
|
|
]
|
|
conditioning, text = run(node, clip, "block1\nBREAK\nblock2")
|
|
|
|
# Concatenated tensor should have seq_len = 10 + 5 = 15
|
|
assert conditioning[0][0].shape[1] == 15
|
|
assert text == "block1\n---\nblock2"
|
|
|
|
|
|
def test_encode_output_text_format(node, mock_clip):
|
|
mock_clip.encode_from_tokens_scheduled.side_effect = [
|
|
[[torch.zeros(1, 10, 768), {"pooled_output": torch.zeros(1, 768)}]],
|
|
[[torch.zeros(1, 5, 768), {"pooled_output": torch.zeros(1, 768)}]],
|
|
[[torch.zeros(1, 8, 768), {"pooled_output": torch.zeros(1, 768)}]],
|
|
]
|
|
_, text = run(node, mock_clip, "a\nBREAK\nb\nBREAK\nc")
|
|
assert text == "a\n---\nb\n---\nc"
|