fixup! Add user interface for text generation

tests/generate/test_integration_transfomers.py

@@ -244,7 +244,26 @@ class Spam(BaseModel):
     assert isinstance(result.bar, float)
     assert isinstance(result.spam, str)
     assert isinstance(result.fuzz, bool)
-    assert len(result.spam) == 10
+    assert len(result.spam) <= 10
+
+
+def test_transformers_json_batch():
+    model_name = "hf-internal-testing/tiny-random-GPTJForCausalLM"
+    model = models.transformers(model_name, device="cpu")
+    prompts = ["Output some JSON ", "Output more JSON"]
+
+    class Spam(BaseModel):
+        foo: int
+        bar: float
+        spam: constr(max_length=10)
+        fuzz: bool
+
+    rng = torch.Generator()
+    rng.manual_seed(0)  # make sure that `bar` is not an int
+
+    result = generate.json(model, Spam, max_tokens=500)(prompts, rng=rng)
+    assert isinstance(result[0], BaseModel)
+    assert isinstance(result[1], BaseModel)
 
 
 def test_transformers_json_str_enum():