revert return dict and wrap apply_chat_template (#691)

Author: awni

README.md

@@ -71,7 +71,7 @@ prompt = "Write a story about Einstein"
 
 messages = [{"role": "user", "content": prompt}]
 prompt = tokenizer.apply_chat_template(
-    messages, add_generation_prompt=True, return_dict=False,
+    messages, add_generation_prompt=True,
 )
 
 text = generate(model, tokenizer, prompt=prompt, verbose=True)
@@ -130,7 +130,7 @@ prompt = "Write a story about Einstein"
 
 messages = [{"role": "user", "content": prompt}]
 prompt = tokenizer.apply_chat_template(
-    messages, add_generation_prompt=True, return_dict=False,
+    messages, add_generation_prompt=True,
 )
 
 for response in stream_generate(model, tokenizer, prompt, max_tokens=512):

mlx_lm/cache_prompt.py

@@ -117,7 +117,6 @@ def main():
             messages,
             add_generation_prompt=False,
             continue_final_message=True,
-            return_dict=False,
         )
 
     else:

mlx_lm/chat.py

@@ -140,7 +140,8 @@ def print_help():
             messages.append({"role": "system", "content": args.system_prompt})
         messages.append({"role": "user", "content": query})
         prompt = tokenizer.apply_chat_template(
-            messages, add_generation_prompt=True, return_dict=False
+            messages,
+            add_generation_prompt=True,
         )
         for response in stream_generate(
             model,

mlx_lm/evaluate.py

@@ -63,7 +63,6 @@ def apply_chat_template(self, chat_history, add_generation_prompt=True) -> str:
             tokenize=False,
             add_generation_prompt=add_generation_prompt,
             continue_final_message=not add_generation_prompt,
-            return_dict=False,
             **extra_kwargs,
         )
 

mlx_lm/examples/batch_generate_response.py

@@ -21,7 +21,6 @@
     tokenizer.apply_chat_template(
         [{"role": "user", "content": p}],
         add_generation_prompt=True,
-        return_dict=False,
     )
     for p in prompts
 ]
@@ -42,7 +41,6 @@
     tokenizer.apply_chat_template(
         [{"role": "user", "content": p}],
         add_generation_prompt=True,
-        return_dict=False,
     )
     for p in prompts
 ]

mlx_lm/examples/chat.py

@@ -16,7 +16,8 @@
 prompt = "Hi my name is <Name>."
 messages = [{"role": "user", "content": prompt}]
 prompt = tokenizer.apply_chat_template(
-    messages, add_generation_prompt=True, return_dict=False
+    messages,
+    add_generation_prompt=True,
 )
 
 # Assistant response
@@ -32,7 +33,8 @@
 prompt = "What's my name?"
 messages = [{"role": "user", "content": prompt}]
 prompt = tokenizer.apply_chat_template(
-    messages, add_generation_prompt=True, return_dict=False
+    messages,
+    add_generation_prompt=True,
 )
 
 # Assistant response

mlx_lm/examples/generate_response.py

@@ -16,7 +16,6 @@
 prompt = tokenizer.apply_chat_template(
     conversation=conversation,
     add_generation_prompt=True,
-    return_dict=False,
 )
 
 # Specify the maximum number of tokens

mlx_lm/examples/sharded_generate.py

@@ -64,7 +64,8 @@ def rprint(*args, **kwargs):
 
     messages = [{"role": "user", "content": args.prompt}]
     prompt = tokenizer.apply_chat_template(
-        messages, add_generation_prompt=True, return_dict=False
+        messages,
+        add_generation_prompt=True,
     )
 
     for response in stream_generate(

mlx_lm/examples/tool_use.py

@@ -34,7 +34,6 @@ def multiply(a: float, b: float):
     messages,
     add_generation_prompt=True,
     tools=list(tools.values()),
-    return_dict=False,
 )
 
 prompt_cache = make_prompt_cache(model)
@@ -63,7 +62,6 @@ def multiply(a: float, b: float):
 prompt = tokenizer.apply_chat_template(
     messages,
     add_generation_prompt=True,
-    return_dict=False,
 )
 
 # Generate the final response:

mlx_lm/generate.py

@@ -1325,7 +1325,6 @@ def main():
             tokenize=False,
             continue_final_message=has_prefill,
             add_generation_prompt=not has_prefill,
-            return_dict=False,
             **template_kwargs,
         )
 
@@ -1337,7 +1336,6 @@ def main():
                 messages,
                 tokenize=False,
                 continue_final_message=has_prefill,
-                return_dict=False,
                 add_generation_prompt=not has_prefill,
             )
             prompt = prompt[test_prompt.index("<query>") :]