feat: implement batch tokenization for TokenizeManager

  - Use tokenizer() for batch encoding plain texts
  - Use apply_chat_template() for batch processing chat templates
  - Remove padding tokens using attention mask
  - Preserve original message order
  - Add comprehensive unit tests for batch tokenization

Author: Alise-svg

python/minisgl/tokenizer/tokenize.py

@@ -7,25 +7,53 @@
 from transformers import PreTrainedTokenizerBase
 
 
-class TokenizeManager:
+class TokenizeManager: 
     def __init__(self, tokenizer: PreTrainedTokenizerBase) -> None:
         self.tokenizer = tokenizer
 
     def tokenize(self, msgs: List[TokenizeMsg]) -> List[torch.Tensor]:
-        results: List[torch.Tensor] = []
-        # TODO: batch tokenization
-        for msg in msgs:
+        if not msgs:
+            return []
+
+        # Separate plain text and chat template messages while preserving order
+        plain_indices: List[int] = []
+        plain_texts: List[str] = []
+        chat_indices: List[int] = []
+        chat_convs: List[List[dict]] = []
+
+        for i, msg in enumerate(msgs):
             if isinstance(msg.text, list):
-                prompt = self.tokenizer.apply_chat_template(
-                    msg.text,
-                    tokenize=False,
-                    add_generation_prompt=True,
-                )
-                assert isinstance(prompt, str)
+                chat_indices.append(i)
+                chat_convs.append(msg.text)
             else:
-                prompt = msg.text
-            input_ids: torch.Tensor = (  # type: ignore
-                self.tokenizer.encode(prompt, return_tensors="pt")
+                plain_indices.append(i)
+                plain_texts.append(msg.text)
+
+        results: List[torch.Tensor | None] = [None] * len(msgs)
+
+        # Batch encode plain texts
+        if plain_texts:
+            encoded = self.tokenizer(plain_texts, return_tensors="pt", padding=True)
+            input_ids = encoded["input_ids"]
+            attention_mask = encoded["attention_mask"]
+            for i, (ids, mask) in enumerate(zip(input_ids, attention_mask)):
+                # Remove padding tokens
+                length = mask.sum().item()
+                results[plain_indices[i]] = ids[:length].to(torch.int32)
+
+        # Batch encode chat templates
+        if chat_convs:
+            prompts = self.tokenizer.apply_chat_template(
+                chat_convs,
+                tokenize=False,
+                add_generation_prompt=True,
             )
-            results.append(input_ids.view(-1).to(torch.int32))
-        return results
+            encoded = self.tokenizer(prompts, return_tensors="pt", padding=True)
+            input_ids = encoded["input_ids"]
+            attention_mask = encoded["attention_mask"]
+            for i, (ids, mask) in enumerate(zip(input_ids, attention_mask)):
+                # Remove padding tokens
+                length = mask.sum().item()
+                results[chat_indices[i]] = ids[:length].to(torch.int32)
+
+        return results  # type: ignore

tests/tokenizer/test_tokenize.py

@@ -0,0 +1,226 @@
+"""Unit tests for TokenizeManager batch tokenization.""" 
+from __future__ import annotations
+
+import torch
+from minisgl.core import SamplingParams
+from minisgl.message import TokenizeMsg
+from minisgl.tokenizer.tokenize import TokenizeManager
+from transformers import AutoTokenizer
+
+
+def get_test_tokenizer():
+    """Get a small tokenizer for testing."""
+    return AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B", trust_remote_code=True)
+
+
+def test_single_plain_text():
+    """Test tokenization of a single plain text message."""
+    tokenizer = get_test_tokenizer()
+    manager = TokenizeManager(tokenizer)
+
+    msg = TokenizeMsg(
+        uid=0,
+        text="Hello, world!",
+        sampling_params=SamplingParams(),
+    )
+
+    results = manager.tokenize([msg])
+
+    assert len(results) == 1
+    expected = tokenizer.encode("Hello, world!", return_tensors="pt").view(-1)
+    assert torch.equal(results[0], expected.to(torch.int32))
+
+
+def test_batch_plain_text():
+    """Test batch tokenization of multiple plain text messages."""
+    tokenizer = get_test_tokenizer()
+    manager = TokenizeManager(tokenizer)
+
+    texts = [
+        "Hello, world!",
+        "How are you?",
+        "This is a test.",
+        "Mini-SGLang is awesome!",
+    ]
+
+    msgs = [
+        TokenizeMsg(uid=i, text=text, sampling_params=SamplingParams())
+        for i, text in enumerate(texts)
+    ]
+
+    results = manager.tokenize(msgs)
+
+    assert len(results) == len(texts)
+
+    for i, (result, text) in enumerate(zip(results, texts)):
+        expected = tokenizer.encode(text, return_tensors="pt").view(-1)
+        assert torch.equal(result, expected.to(torch.int32)), f"Mismatch at index {i}"
+
+
+def test_single_chat_template():
+    """Test tokenization of a single chat template message."""
+    tokenizer = get_test_tokenizer()
+    manager = TokenizeManager(tokenizer)
+
+    msg = TokenizeMsg(
+        uid=0,
+        text=[{"role": "user", "content": "Hello!"}],
+        sampling_params=SamplingParams(),
+    )
+
+    results = manager.tokenize([msg])
+
+    assert len(results) == 1
+
+    # Verify the result is valid tokens
+    assert results[0].dtype == torch.int32
+    assert len(results[0]) > 0
+
+
+def test_batch_chat_template():
+    """Test batch tokenization of multiple chat template messages."""
+    tokenizer = get_test_tokenizer()
+    manager = TokenizeManager(tokenizer)
+
+    conversations = [
+        [{"role": "user", "content": "Hello!"}],
+        [{"role": "user", "content": "How are you?"}],
+        [
+            {"role": "system", "content": "You are helpful."},
+            {"role": "user", "content": "Hi!"},
+        ],
+    ]
+
+    msgs = [
+        TokenizeMsg(uid=i, text=conv, sampling_params=SamplingParams())
+        for i, conv in enumerate(conversations)
+    ]
+
+    results = manager.tokenize(msgs)
+
+    assert len(results) == len(conversations)
+
+    for result in results:
+        assert result.dtype == torch.int32
+        assert len(result) > 0
+
+
+def test_mixed_batch():
+    """Test batch tokenization with mixed plain text and chat template messages."""
+    tokenizer = get_test_tokenizer()
+    manager = TokenizeManager(tokenizer)
+
+    msgs = [
+        TokenizeMsg(uid=0, text="Plain text message", sampling_params=SamplingParams()),
+        TokenizeMsg(
+            uid=1, text=[{"role": "user", "content": "Chat message"}], sampling_params=SamplingParams()
+        ),
+        TokenizeMsg(uid=2, text="Another plain text", sampling_params=SamplingParams()),
+        TokenizeMsg(
+            uid=3,
+            text=[
+                {"role": "system", "content": "System prompt"},
+                {"role": "user", "content": "User message"},
+            ],
+            sampling_params=SamplingParams(),
+        ),
+    ]
+
+    results = manager.tokenize(msgs)
+
+    assert len(results) == 4
+
+    # Verify plain text results
+    expected_0 = tokenizer.encode("Plain text message", return_tensors="pt").view(-1)
+    assert torch.equal(results[0], expected_0.to(torch.int32))
+
+    expected_2 = tokenizer.encode("Another plain text", return_tensors="pt").view(-1)
+    assert torch.equal(results[2], expected_2.to(torch.int32))
+
+    # Verify chat template results are valid
+    assert results[1].dtype == torch.int32
+    assert results[3].dtype == torch.int32
+
+
+def test_empty_batch():
+    """Test tokenization of an empty batch."""
+    tokenizer = get_test_tokenizer()
+    manager = TokenizeManager(tokenizer)
+
+    results = manager.tokenize([])
+
+    assert len(results) == 0
+
+
+def test_output_dtype():
+    """Verify that output tensors are int32 as expected by the system."""
+    tokenizer = get_test_tokenizer()
+    manager = TokenizeManager(tokenizer)
+
+    msgs = [
+        TokenizeMsg(uid=0, text="Test", sampling_params=SamplingParams()),
+        TokenizeMsg(uid=1, text=[{"role": "user", "content": "Test"}], sampling_params=SamplingParams()),
+    ]
+
+    results = manager.tokenize(msgs)
+
+    for result in results:
+        assert result.dtype == torch.int32, f"Expected int32, got {result.dtype}"
+
+
+def test_consistency_with_original():
+    """Verify batch tokenization produces same results as individual tokenization."""
+    tokenizer = get_test_tokenizer()
+    manager = TokenizeManager(tokenizer)
+
+    texts = ["First message", "Second message", "Third message"]
+
+    msgs = [
+        TokenizeMsg(uid=i, text=text, sampling_params=SamplingParams())
+        for i, text in enumerate(texts)
+    ]
+
+    # Batch tokenization
+    batch_results = manager.tokenize(msgs)
+
+    # Individual tokenization (original behavior)
+    individual_results = []
+    for text in texts:
+        ids = tokenizer.encode(text, return_tensors="pt").view(-1).to(torch.int32)
+        individual_results.append(ids)
+
+    # Compare
+    for i, (batch, individual) in enumerate(zip(batch_results, individual_results)):
+        assert torch.equal(batch, individual), f"Mismatch at index {i}"
+
+
+if __name__ == "__main__":
+    import sys
+
+    failed = False
+    tests = [
+        test_single_plain_text,
+        test_batch_plain_text,
+        test_single_chat_template,
+        test_batch_chat_template,
+        test_mixed_batch,
+        test_empty_batch,
+        test_output_dtype,
+        test_consistency_with_original,
+    ]
+
+    for test in tests:
+        try:
+            test()
+            print(f"✓ {test.__name__}")
+        except AssertionError as e:
+            print(f"✗ {test.__name__}: {e}")
+            failed = True
+        except Exception as e:
+            print(f"✗ {test.__name__}: {type(e).__name__}: {e}")
+            failed = True
+
+    if failed:
+        sys.exit(1)
+    else:
+        print("\nAll tests passed!")