java-topology/defects/transformers-0004/test/Transformers0004Wav2Vec2EsmSpecialIdsTest.py

"""
transformers-0004: wav2vec2, wav2vec2_phoneme, esm - all_special_ids property rebuilt per loop iteration.
CWE-407 / MOAD-0001 - The Sedimentary Defect

Stubs PreTrainedTokenizer-style property behavior to confirm:
1. The @property all_special_ids/all_special_tokens is accessed once per element in a loop
   (reproduces the defect), and
2. The patched code (set cached before loop) accesses it exactly once regardless of T.

No model download required. Pure data structure test.
"""

import unittest


class MockPropertyCounter:
    """Mimics a @property that rebuilds a list each call, tracking access count."""

    def __init__(self, items):
        self._items = items
        self.access_count = 0

    @property
    def all_special_ids(self):
        self.access_count += 1
        return list(self._items)  # rebuild each access like the real property

    @property
    def all_special_tokens(self):
        self.access_count += 1
        return [str(i) for i in self._items]  # rebuild each access


class DefectiveDecode:
    """Reproduces the wav2vec2 defect: all_special_ids accessed per token."""

    def __init__(self, prop_counter):
        self._prop = prop_counter

    def convert_ids_to_tokens_buggy(self, ids, skip_special_tokens=True):
        tokens = []
        for index in ids:
            # DEFECT: self.all_special_ids rebuilt each iteration
            if skip_special_tokens and index in self._prop.all_special_ids:
                continue
            tokens.append(str(index))
        return tokens


class PatchedDecode:
    """Patched version: cache set once before loop."""

    def __init__(self, prop_counter):
        self._prop = prop_counter

    def convert_ids_to_tokens_fixed(self, ids, skip_special_tokens=True):
        tokens = []
        special_ids = set(self._prop.all_special_ids) if skip_special_tokens else None
        for index in ids:
            if special_ids is not None and index in special_ids:
                continue
            tokens.append(str(index))
        return tokens


class DefectiveEsmMask:
    """Reproduces the ESM defect: all_special_ids accessed per element in list comprehension."""

    def __init__(self, prop_counter):
        self._prop = prop_counter

    def get_special_tokens_mask_buggy(self, token_ids_0):
        # DEFECT: self.all_special_ids rebuilt per element
        return [1 if token in self._prop.all_special_ids else 0 for token in token_ids_0]


class PatchedEsmMask:
    """Patched ESM: cache set once before comprehension."""

    def __init__(self, prop_counter):
        self._prop = prop_counter

    def get_special_tokens_mask_fixed(self, token_ids_0):
        special_ids = set(self._prop.all_special_ids)
        return [1 if token in special_ids else 0 for token in token_ids_0]


class Transformers0004Wav2Vec2EsmSpecialIdsTest(unittest.TestCase):

    # Special token IDs: 0=pad, 1=bos, 2=eos, 3=unk (S=4)
    SPECIAL_IDS = {0, 1, 2, 3}
    T = 64  # sequence length for test

    def _make_sequence(self):
        """Token sequence with some special tokens scattered throughout."""
        ids = []
        for i in range(self.T):
            # Every 8th token is special
            ids.append(i % max(self.SPECIAL_IDS) if i % 8 == 0 else 10 + i)
        return ids

    # --- MOAD-0001 reproduce: property access count grows with T ---

    def test_defect_wav2vec2_property_accessed_once_per_token(self):
        """DEFECT: all_special_ids accessed T times (once per token in loop)."""
        ids = self._make_sequence()
        counter = MockPropertyCounter(self.SPECIAL_IDS)
        decoder = DefectiveDecode(counter)
        decoder.convert_ids_to_tokens_buggy(ids, skip_special_tokens=True)
        self.assertEqual(
            counter.access_count,
            len(ids),
            f"Defective code accessed all_special_ids {counter.access_count}x "
            f"(expected {len(ids)}, one per token)"
        )

    def test_patch_wav2vec2_property_accessed_exactly_once(self):
        """PATCH: all_special_ids accessed exactly once regardless of T."""
        ids = self._make_sequence()
        counter = MockPropertyCounter(self.SPECIAL_IDS)
        decoder = PatchedDecode(counter)
        decoder.convert_ids_to_tokens_fixed(ids, skip_special_tokens=True)
        self.assertEqual(
            counter.access_count,
            1,
            f"Patched code accessed all_special_ids {counter.access_count}x "
            f"(expected exactly 1)"
        )

    def test_defect_esm_property_accessed_once_per_element(self):
        """DEFECT: ESM all_special_ids accessed T times in list comprehension."""
        ids = self._make_sequence()
        counter = MockPropertyCounter(self.SPECIAL_IDS)
        masker = DefectiveEsmMask(counter)
        masker.get_special_tokens_mask_buggy(ids)
        self.assertEqual(
            counter.access_count,
            len(ids),
            f"Defective ESM code accessed all_special_ids {counter.access_count}x "
            f"(expected {len(ids)})"
        )

    def test_patch_esm_property_accessed_exactly_once(self):
        """PATCH: ESM all_special_ids accessed exactly once."""
        ids = self._make_sequence()
        counter = MockPropertyCounter(self.SPECIAL_IDS)
        masker = PatchedEsmMask(counter)
        masker.get_special_tokens_mask_fixed(ids)
        self.assertEqual(
            counter.access_count,
            1,
            f"Patched ESM code accessed all_special_ids {counter.access_count}x "
            f"(expected exactly 1)"
        )

    # --- Correctness: patch produces same results as defective code ---

    def test_wav2vec2_correctness(self):
        """Patched decode produces same token list as defective version."""
        ids = self._make_sequence()
        counter_buggy = MockPropertyCounter(self.SPECIAL_IDS)
        counter_fixed = MockPropertyCounter(self.SPECIAL_IDS)
        buggy = DefectiveDecode(counter_buggy).convert_ids_to_tokens_buggy(ids)
        fixed = PatchedDecode(counter_fixed).convert_ids_to_tokens_fixed(ids)
        self.assertEqual(buggy, fixed, "Patched and defective code produce different results")

    def test_esm_mask_correctness(self):
        """Patched ESM mask matches defective version element-by-element."""
        ids = self._make_sequence()
        counter_buggy = MockPropertyCounter(self.SPECIAL_IDS)
        counter_fixed = MockPropertyCounter(self.SPECIAL_IDS)
        buggy_mask = DefectiveEsmMask(counter_buggy).get_special_tokens_mask_buggy(ids)
        fixed_mask = PatchedEsmMask(counter_fixed).get_special_tokens_mask_fixed(ids)
        self.assertEqual(buggy_mask, fixed_mask, "Patched and defective masks differ")

    # --- Speedup ratio validation ---

    def test_speedup_ratio(self):
        """Property access count ratio must be >= T for the defective case."""
        ids = list(range(self.T))
        counter = MockPropertyCounter(self.SPECIAL_IDS)
        DefectiveDecode(counter).convert_ids_to_tokens_buggy(ids)
        ratio = counter.access_count  # should equal T
        self.assertGreaterEqual(
            ratio,
            self.T,
            f"Expected ratio >= {self.T}, got {ratio}"
        )

    def test_wav2vec2_skip_false_no_property_access(self):
        """When skip_special_tokens=False, all_special_ids must not be accessed."""
        ids = self._make_sequence()
        counter = MockPropertyCounter(self.SPECIAL_IDS)
        PatchedDecode(counter).convert_ids_to_tokens_fixed(ids, skip_special_tokens=False)
        self.assertEqual(
            counter.access_count,
            0,
            "Property should not be accessed when skip_special_tokens=False"
        )

    def test_wav2vec2_phoneme_type_fix(self):
        """
        wav2vec2_phoneme DEFECT: compared str tokens against list[int] (all_special_ids),
        always False. Patch uses all_special_tokens (list[str]) correctly.
        """
        special_ids = {0, 1, 2}
        # Token 0 as str (from convert_ids_to_tokens) vs int set: type mismatch
        str_token = "0"
        self.assertNotIn(str_token, list(special_ids),
                         "str token correctly not in int list (type mismatch)")
        # Patched: use str representation
        special_tokens = set(str(i) for i in special_ids)
        self.assertIn(str_token, special_tokens,
                      "str token correctly found in str set (patched)")


if __name__ == "__main__":
    unittest.main(verbosity=2)