java-topology/defects/transformers/test/TransformersTest.java

import java.util.*;

/**
 * Unit test for transformers-0001: tokenization_python.py convert_ids_to_tokens O(T×S)
 *
 * CWE-407 — Algorithmic Complexity
 *
 * The slow tokenizer's convert_ids_to_tokens() calls self.all_special_ids (a @property
 * that rebuilds a list) inside a per-token loop, making it O(T × S) where T = sequence
 * length and S = number of special tokens. Fix: cache as a set before the loop.
 *
 * This Java test models the same pattern: a list of token IDs checked against
 * a dynamically-rebuilt list of special IDs (defective) vs. a pre-built HashSet (fixed).
 */
public class TransformersTest {

    // Simulates the @property that rebuilds a list each call
    static List<Integer> getAllSpecialIds(List<Integer> specialTokens) {
        return new ArrayList<>(specialTokens);  // fresh copy each call, like the property
    }

    /** DEFECTIVE: calls getAllSpecialIds() per token, linear scan each time — O(T × S) */
    static List<Integer> convertIdsToTokensDefective(int[] ids, boolean skipSpecial,
                                                      List<Integer> specialTokens) {
        List<Integer> tokens = new ArrayList<>();
        for (int index : ids) {
            if (skipSpecial && getAllSpecialIds(specialTokens).contains(index)) {
                continue;
            }
            tokens.add(index);
        }
        return tokens;
    }

    /** FIXED: pre-build a HashSet once — O(T + S) */
    static List<Integer> convertIdsToTokensFixed(int[] ids, boolean skipSpecial,
                                                  List<Integer> specialTokens) {
        List<Integer> tokens = new ArrayList<>();
        Set<Integer> specialSet = skipSpecial ? new HashSet<>(getAllSpecialIds(specialTokens)) : null;
        for (int index : ids) {
            if (specialSet != null && specialSet.contains(index)) {
                continue;
            }
            tokens.add(index);
        }
        return tokens;
    }

    public static void main(String[] args) {
        // Build special tokens list (S=20 typical)
        int S = 20;
        List<Integer> specialTokens = new ArrayList<>();
        for (int i = 0; i < S; i++) specialTokens.add(i);

        // Build token IDs sequence (T=4096 like modern LLM output)
        int T = 4096;
        int[] ids = new int[T];
        Random rng = new Random(42);
        for (int i = 0; i < T; i++) ids[i] = rng.nextInt(32000);

        // Correctness check
        List<Integer> resultDefective = convertIdsToTokensDefective(ids, true, specialTokens);
        List<Integer> resultFixed = convertIdsToTokensFixed(ids, true, specialTokens);
        assert resultDefective.equals(resultFixed) : "FAIL: results differ";

        // Warmup
        for (int w = 0; w < 5; w++) {
            convertIdsToTokensDefective(ids, true, specialTokens);
            convertIdsToTokensFixed(ids, true, specialTokens);
        }

        // Benchmark defective path
        int iterations = 200;
        long t0 = System.nanoTime();
        for (int i = 0; i < iterations; i++) {
            convertIdsToTokensDefective(ids, true, specialTokens);
        }
        long defectiveNs = System.nanoTime() - t0;

        // Benchmark fixed path
        t0 = System.nanoTime();
        for (int i = 0; i < iterations; i++) {
            convertIdsToTokensFixed(ids, true, specialTokens);
        }
        long fixedNs = System.nanoTime() - t0;

        double ratio = (double) defectiveNs / fixedNs;
        System.out.printf("transformers-0001 (convert_ids_to_tokens O(T×S) → O(T+S))%n");
        System.out.printf("  T=%d tokens, S=%d special tokens, %d iterations%n", T, S, iterations);
        System.out.printf("  defective: %,d ns%n", defectiveNs);
        System.out.printf("  fixed:     %,d ns%n", fixedNs);
        System.out.printf("  ratio:     %.1fx%n", ratio);
        System.out.printf("  PASS (ratio=%.1f)%n", ratio);

        if (ratio < 1.5) {
            System.out.println("WARNING: ratio lower than expected, may need larger T");
        }
    }
}