Fix generation input handling for chat template

app.py

@@ -41,6 +41,16 @@ EXAMPLES = [
 ]
 
 
+def _build_inputs(tokenizer: AutoTokenizer, prompt: str, device: torch.device):
+    messages = [{"role": "user", "content": prompt}]
+    chat = tokenizer.apply_chat_template(
+        messages,
+        add_generation_prompt=True,
+        tokenize=False,
+    )
+    return tokenizer(chat, return_tensors="pt").to(device)
+
+
 def _generate_one(
     key: str,
     prompt: str,
@@ -50,11 +60,8 @@ def _generate_one(
 ) -> tuple[str, float]:
     tokenizer = tokenizers[key]
     model = models[key]
-    inputs = tokenizer.apply_chat_template(
-        [{"role": "user", "content": prompt}],
-        add_generation_prompt=True,
-        return_tensors="pt",
-    ).to(model.device)
+    inputs = _build_inputs(tokenizer, prompt, model.device)
+    input_ids = inputs["input_ids"]
 
     gen_kwargs: dict = {
         "max_new_tokens": max_new_tokens,
@@ -69,11 +76,11 @@ def _generate_one(
 
     started = time.perf_counter()
     with torch.inference_mode():
-        output = model.generate(inputs, **gen_kwargs)
+        output = model.generate(**inputs, **gen_kwargs)
     elapsed = time.perf_counter() - started
 
     response = tokenizer.decode(
-        output[0, inputs.shape[-1] :],
+        output[0, input_ids.shape[-1] :],
         skip_special_tokens=True,
     ).strip()
     return response, elapsed