Explicitly set return_dict for apply_chat_template (#33372)

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

examples/offline_inference/prompt_embed_inference.py

@@ -38,8 +38,8 @@ def get_prompt_embeds(
     embedding_layer: torch.nn.Module,
 ):
     token_ids = tokenizer.apply_chat_template(
-        chat, add_generation_prompt=True, return_tensors="pt"
-    )
+        chat, add_generation_prompt=True, return_tensors="pt", return_dict=True
+    ).input_ids
     prompt_embeds = embedding_layer(token_ids).squeeze(0)
     return prompt_embeds
 

examples/online_serving/prompt_embed_inference_with_openai_client.py

@@ -49,8 +49,8 @@ def main():
     # Refer to the HuggingFace repo for the correct format to use
     chat = [{"role": "user", "content": "Please tell me about the capital of France."}]
     token_ids = tokenizer.apply_chat_template(
-        chat, add_generation_prompt=True, return_tensors="pt"
-    )
+        chat, add_generation_prompt=True, return_tensors="pt", return_dict=True
+    ).input_ids
 
     embedding_layer = transformers_model.get_input_embeddings()
     prompt_embeds = embedding_layer(token_ids).squeeze(0)

examples/online_serving/token_generation_client.py

@@ -27,7 +27,8 @@ def main(client):
         messages,
         add_generation_prompt=True,
         enable_thinking=False,
-    )
+        return_dict=True,
+    ).input_ids
     payload = {
         "model": MODEL_NAME,
         "token_ids": token_ids,