Handle Transformers BatchEncoding in ZeroGPU runtime

figment/zerogpu_runtime.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+from collections.abc import Mapping
 import json
 import os
 import threading
@@ -91,17 +92,26 @@ class _ZeroGpuRuntime:
     def generate_json(self, prompt: str) -> dict[str, Any]:
         messages = [{"role": "user", "content": prompt}]
         with self.lock:
-            input_ids = self.tokenizer.apply_chat_template(
+            encoded = self.tokenizer.apply_chat_template(
                 messages,
                 add_generation_prompt=True,
                 return_tensors="pt",
             )
             device = next(self.model.parameters()).device
+            if isinstance(encoded, Mapping):
+                input_ids = encoded["input_ids"]
+                attention_mask = encoded.get("attention_mask")
+            else:
+                input_ids = encoded
+                attention_mask = None
             input_ids = input_ids.to(device)
+            if attention_mask is None:
+                attention_mask = self.torch.ones_like(input_ids)
+            else:
+                attention_mask = attention_mask.to(device)
             input_len = int(input_ids.shape[-1])
             available_tokens = max(1, self.max_context_tokens - input_len - 8)
             max_new_tokens = max(1, min(self.max_generation_tokens, available_tokens))
-            attention_mask = self.torch.ones_like(input_ids)
             with self.torch.inference_mode():
                 output_ids = self.model.generate(
                     input_ids=input_ids,

tests/test_zerogpu_runtime.py

@@ -62,3 +62,84 @@ def test_zerogpu_runtime_uses_native_transformers_with_mamba_kernels_disabled(mo
     assert calls["model"]["kwargs"]["torch_dtype"] == "bf16"
     assert calls["device"] == "cuda"
     assert calls["eval"] is True
+
+
+def test_zerogpu_runtime_generate_json_accepts_batch_encoding_tokenizer_output() -> None:
+    calls: dict[str, Any] = {}
+
+    class FakeTensor:
+        def __init__(self, shape: tuple[int, int]) -> None:
+            self.shape = shape
+
+        def to(self, device: str) -> "FakeTensor":
+            calls.setdefault("moved_to", []).append(device)
+            return self
+
+    class FakeGeneratedIds:
+        def __getitem__(self, key: Any) -> list[int]:
+            calls["generated_slice"] = key
+            return [4, 5, 6]
+
+    class FakeTokenizer:
+        pad_token_id = 0
+        eos_token_id = 2
+
+        def apply_chat_template(self, *_: Any, **__: Any) -> dict[str, FakeTensor]:
+            return {
+                "input_ids": FakeTensor((1, 3)),
+                "attention_mask": FakeTensor((1, 3)),
+            }
+
+        def decode(self, generated_ids: list[int], *, skip_special_tokens: bool) -> str:
+            calls["decoded"] = {"ids": generated_ids, "skip_special_tokens": skip_special_tokens}
+            return '{"protocol_urgency": "monitor", "source_cards": []}'
+
+    class FakeTorch:
+        @staticmethod
+        def ones_like(tensor: FakeTensor) -> FakeTensor:
+            calls["ones_like"] = tensor
+            return FakeTensor(tensor.shape)
+
+        @staticmethod
+        def inference_mode() -> Any:
+            class Context:
+                def __enter__(self) -> None:
+                    return None
+
+                def __exit__(self, *_: Any) -> None:
+                    return None
+
+            return Context()
+
+    class FakeModel:
+        def parameters(self) -> Any:
+            return iter([SimpleNamespace(device="cuda")])
+
+        def generate(self, **kwargs: Any) -> FakeGeneratedIds:
+            calls["generate"] = kwargs
+            return FakeGeneratedIds()
+
+    class FakeLock:
+        def __enter__(self) -> None:
+            return None
+
+        def __exit__(self, *_: Any) -> None:
+            return None
+
+    runtime = object.__new__(_ZeroGpuRuntime)
+    runtime.model_id = "model"
+    runtime.max_context_tokens = 32
+    runtime.max_generation_tokens = 8
+    runtime.lock = FakeLock()
+    runtime.torch = FakeTorch()
+    runtime.tokenizer = FakeTokenizer()
+    runtime.model = FakeModel()
+
+    result = runtime.generate_json("Return JSON.")
+
+    assert result == {"protocol_urgency": "monitor", "source_cards": []}
+    assert calls["moved_to"] == ["cuda", "cuda"]
+    assert calls["generate"]["input_ids"].shape == (1, 3)
+    assert calls["generate"]["attention_mask"].shape == (1, 3)
+    assert calls["generated_slice"] == (0, slice(3, None, None))
+    assert calls["decoded"] == {"ids": [4, 5, 6], "skip_special_tokens": True}