Update README.md

README.md

@@ -50,17 +50,20 @@ import json
 from typing import Any, Dict, List
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
+# Specify the desired model name here
 model_name = "katanemo/Arch-Agent-3B"
+
 model = AutoModelForCausalLM.from_pretrained(
     model_name, device_map="auto", torch_dtype="auto", trust_remote_code=True
 )
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-# Please use our provided prompt for best performance
+
+# Please use the recommended prompt for each model.
 TASK_PROMPT = (
     "You are a helpful assistant designed to assist with the user query by making one or more function calls if needed."
     "\n\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\n"
-    "You are provided with function signatures within <tools></tools> XML tags:\n<tools>{tool_text}"
+    "You are provided with function signatures within <tools></tools> XML tags:\n<tools>\n{tool_text}"
     "\n</tools>\n\nFor each function call, return a json object with function name and arguments within "
     """<tool_call></tool_call> XML tags:\n<tool_call>\n{{"name": <function-name>, """
     """"arguments": <args-json-object>}}\n</tool_call>"""
@@ -108,12 +111,12 @@ messages = [
     {"role": "user", "content": "What is the weather in Seattle?"},
 ]
 
+#### 2.2.3 Run inference
 model_inputs = tokenizer.apply_chat_template(
-    messages, add_generation_prompt=True, return_tensors="pt"
+    messages, add_generation_prompt=True, return_tensors="pt", return_dict=True
 ).to(model.device)
 
 generated_ids = model.generate(**model_inputs, max_new_tokens=32768)
-
 generated_ids = [
     output_ids[len(input_ids) :]
     for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)