From b790a7ff290f64f865e0c66602a98ba70a901578 Mon Sep 17 00:00:00 2001
From: Don Mahurin <2797413+dmahurin@users.noreply.github.com>
Date: Sat, 28 Sep 2024 14:10:55 -0700
Subject: [PATCH] Allow simpler function calling sytax, like used with Phi-3
 function calling model

---
 examples/function-calling/README.md           |  5 +++
 examples/function-calling/function_tool.py    | 35 ++++++++++++++-----
 .../llama-cli-function-runner.py              |  4 +--
 3 files changed, 33 insertions(+), 11 deletions(-)

diff --git a/examples/function-calling/README.md b/examples/function-calling/README.md
index f431c9835..e245747b1 100644
--- a/examples/function-calling/README.md
+++ b/examples/function-calling/README.md
@@ -44,3 +44,8 @@ What is 37234 times 39?
 To calculate 37234 times 39, I'll perform the multiplication. Let's do that.
 The result of multiplying 37234 by 39 is 1,452,126. If you have any more calculations or questions, feel free to ask!
 ```
+
+## Function calling example, using Phi-3 function calling
+```
+./examples/function-calling/llama-cli-function-runner.py -m `huggingface-cli download nold/Phi-3-mini-4k-instruct-function-calling-GGUF Phi-3-mini-4k-instruct-function-calling_Q4_K_M.gguf` --special --display-prompt -i
+```
diff --git a/examples/function-calling/function_tool.py b/examples/function-calling/function_tool.py
index 850a1f715..f92d81ec1 100644
--- a/examples/function-calling/function_tool.py
+++ b/examples/function-calling/function_tool.py
@@ -64,6 +64,9 @@ def generate_functionary_schema_from_functions(functions, namespace="functions")
     schema += "}} // namespace {}".format(namespace)
     return schema
 
+def generate_simple_schema_from_functions(functions) -> str:
+    return '\n'.join([json.dumps(function).replace('{', '{ ').replace('}', ' }') for function in functions])
+
 functionary_prompt_start = """<|start_header_id|>system<|end_header_id|>
 
 You are capable of executing available function(s) if required.
@@ -81,13 +84,27 @@ functionary_prompt_end = """<|eot_id|><|start_header_id|>system<|end_header_id|>
 When you send a message containing Python code to python, it will be executed in a stateful Jupyter notebook environment. python will respond with the output of the execution or time out after 60.0 seconds. The drive at '/mnt/data' can be used to save and persist user files.<|eot_id|><|start_header_id|>user<|end_header_id|>
 """
 
+simple_prompt_start = """<s><|user|> You are a helpful assistant with access to the following functions. Use them if required - """
+simple_prompt_end = """<|end|>"""
+
 def get_chat_tool_format(args, tools):
-    return {
-        'prompt': functionary_prompt_start + generate_functionary_schema_from_functions(tools) + functionary_prompt_end,
-        'function_marker': '>>>',
-        'function_re': r'>>>([^\n]*)\n(.*)<\|eot_id\|>',
-        'user_start': '<|start_header_id|>user<|end_header_id|>\n',
-        'user_end': '<|eot_id|><|start_header_id|>assistant<|end_header_id|>' + '\n',
-        'tool_start': '',
-        'tool_end': '<|eot_id|><|start_header_id|>assistant<|end_header_id|>'
-    }
+    if 'functionary' in args.model.lower():
+        return {
+            'prompt': functionary_prompt_start + generate_functionary_schema_from_functions(tools) + functionary_prompt_end,
+            'function_marker': '>>>',
+            'function_re': r'>>>([^\n]*)\n(.*)<\|eot_id\|>',
+            'user_start': '<|start_header_id|>user<|end_header_id|>\n',
+            'user_end': '<|eot_id|><|start_header_id|>assistant<|end_header_id|>' + '\n',
+            'tool_start': '',
+            'tool_end': '<|eot_id|><|start_header_id|>assistant<|end_header_id|>'
+        }
+    else:
+        return {
+            'prompt': simple_prompt_start + generate_simple_schema_from_functions(tools) + simple_prompt_end,
+            'function_marker': '<functioncall>',
+            'function_re': r'<functioncall> \n?(.*)<\|end\|>',
+            'user_start': '<|user|> ',
+            'user_end': '<|end|>' + '\n',
+            'tool_start': '<|user|>',
+            'tool_end': '<|end|> <|assistant|>'
+        }
diff --git a/examples/function-calling/llama-cli-function-runner.py b/examples/function-calling/llama-cli-function-runner.py
index 473139bdb..066363b69 100755
--- a/examples/function-calling/llama-cli-function-runner.py
+++ b/examples/function-calling/llama-cli-function-runner.py
@@ -23,6 +23,7 @@ def main():
     parser.add_argument('--display-prompt', action=argparse.BooleanOptionalAction, default=False)
     parser.add_argument('--special', action=argparse.BooleanOptionalAction, default=False)
     parser.add_argument('--reverse-prompt', type=str)
+    parser.add_argument('-m', '--model', type=str, default='model.gguf')
     parser.add_argument('--ctx-size', type=int, default=1024)
     args, other_args = parser.parse_known_args()
 
@@ -31,8 +32,7 @@ def main():
 
     if args.display_prompt: print(tool_format['prompt'])
 
-    command = [ './llama-cli', '-i', '-p', tool_format['prompt'], '--reverse-prompt', args.reverse_prompt, '--escape', '--special', '--no-display-prompt', '--log-disable', '--simple-io', '--ctx-size',  str(args.ctx_size), *other_args]
-    print("'" + "' '".join(command) + "'")
+    command = [ './llama-cli', '-i', '-p', tool_format['prompt'], '--model', args.model, '--reverse-prompt', args.reverse_prompt, '--escape', '--special', '--no-display-prompt', '--log-disable', '--simple-io', '--ctx-size',  str(args.ctx_size), *other_args]
 
     process = subprocess.Popen(
         command,