ftobler/llama
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

tool functions and pytest

Browse Source
This commit is contained in:
Florin Tobler
2025-01-01 18:20:50 +01:00
parent 823f13ab51
commit fd7e3d5235
10 changed files with 406 additions and 108 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
.gitignore vendored
View File

@@ -1,2 +1,3 @@
/model/*
*.prof
*.prof
__pycache__

8
__main__.py Normal file
View File

@@ -0,0 +1,8 @@
print("running __main__.-py")
from llama import main
if __name__ == "__main__":
main()

228
llama.py
View File

@@ -2,8 +2,9 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
import time
import torch
import random
import datetime
import json
from tool_helper import tool_list, parse_and_execute_tool_call
from tool_functions import register_dummy
import utils
t_start = time.time()
@@ -41,10 +42,9 @@ print("load took %.3fs" % (time.time() - t_start))
max_context_length = model.config.max_position_embeddings
# if tokenizer.chat_template is None:
print("apply external chat template...")
with open("chat_template.json", "r") as f:
tokenizer.chat_template = json.load(f)
tokenizer.chat_template = utils.load_json_file("chat_template.json")
print("max_context_length is %d tokens." % (max_context_length))
@@ -97,22 +97,9 @@ messages = [
roleflip = {"role": "system", "content": "Keep the conversation going, ask for more information on the subject. Keep messages short at max 1-2 sentences. Do not thank and say goodbye."}
def current_time():
"""Get the current local date and time as a string."""
return datetime.datetime.now().strftime("%Y-%m-%d %H:%M")
def random_float():
"""Get a random float from 0..1"""
return str(random.random())
def random_int(a: int, b: int):
"""Return random integer in range [a, b], including both end points.
Args:
a: minimum possible value
b: maximum possible value"""
return str(random.randint(a, b))
tool_functions = [current_time, random_float, random_int]
register_dummy()
# tool_functions = [current_time, random_float, random_int]
@@ -139,31 +126,35 @@ def generate_incremental(inputs):
generated_tokens = input_ids # Initially, this is just the input tokens
n = 0
try:
# Loop to generate one token at a time
while True:
# Call the model with the current tokens
outputs = model(input_ids=generated_tokens, use_cache=True)
# Loop to generate one token at a time
while True:
# Call the model with the current tokens
outputs = model(input_ids=generated_tokens, use_cache=True)
# Get the next token (the last token from the generated sequence)
next_token = outputs.logits.argmax(dim=-1)[:, -1]
# Get the next token (the last token from the generated sequence)
next_token = outputs.logits.argmax(dim=-1)[:, -1]
# Append the new token to the sequence
generated_tokens = torch.cat([generated_tokens, next_token.unsqueeze(0)], dim=1)
# Append the new token to the sequence
generated_tokens = torch.cat([generated_tokens, next_token.unsqueeze(0)], dim=1)
# Decode and print the newly generated token (skip special tokens)
out_text = tokenizer.decode(next_token, skip_special_tokens=True)
print(out_text, end="", flush=True) # Print without newline
# Decode and print the newly generated token (skip special tokens)
out_text = tokenizer.decode(next_token, skip_special_tokens=True)
print(out_text, end="", flush=True) # Print without newline
# Check if the generated token is the end-of-sequence token
if next_token.item() == tokenizer.eos_token_id:
print("")
break
# Check if the generated token is the end-of-sequence token
if next_token.item() == tokenizer.eos_token_id:
print("")
break
n += 1
if n >= 30:
n = 0
torch.cuda.empty_cache()
n += 1
if n >= 15:
n = 0
torch.cuda.empty_cache()
except KeyboardInterrupt:
pass
# Once done, return the full generated sequence
@@ -184,7 +175,7 @@ def append_generate_chat(input_text: str, role="user"):
# input_text = "Hello, who are you?"
# inputs = tokenizer(input_text, return_tensors="pt").to("cpu") # .to("cuda") .to("cpu")
inputs = tokenizer.apply_chat_template(messages, return_tensors="pt", tokenize=True, return_dict=True, add_generation_prompt=True, tools=tool_functions) #continue_final_message=True,
inputs = tokenizer.apply_chat_template(messages, return_tensors="pt", tokenize=True, return_dict=True, add_generation_prompt=True, tools=tool_list) #continue_final_message=True,
inputs = {key: value.to(model.device) for key, value in inputs.items()}
# inputs = {key: value.to("cpu") for key, value in inputs.items()}
# inputs["input_ids"] = inputs["input_ids"][:, 1:]
@@ -194,82 +185,105 @@ def append_generate_chat(input_text: str, role="user"):
# append result to message history
messages.append({"role": "assistant", "content": out_text})
print("")
print("generation took %.3fs (%d tokens)" % (time.time() - t_start, len(outputs[0])))
while True:
# print an input prompt to receive text or commands
input_text = input(">>> ")
print("")
# handle tool call and check if a tool call has happened.
tool_result = parse_and_execute_tool_call(out_text, tool_list)
if tool_result != None:
# tool call happened
# tool_result = "<tool_response>%s</tool_response>" % tool_result
# depending on the chat template the tool response tags must or must not be passed. :(
append_generate_chat(tool_result, role="tool")
if input_text.startswith("!"):
# append_generate_chat("<tool_response>%s</tool_response>" % input_text[1:], role="tool")
append_generate_chat("%s" % input_text[1:], role="tool") # depending on the chat template the tool response tags must or must not be passed. :(
elif input_text.startswith("/clear"):
print("clearing chat history")
messages = [messages[0]]
def main():
global messages
while True:
# print an input prompt to receive text or commands
input_text = input(">>> ")
print("")
elif input_text.startswith("/history"):
history = tokenizer.apply_chat_template(messages, return_tensors="pt", tokenize=False, add_generation_prompt=False, tools=tool_functions)
print(history)
elif input_text.startswith("/undo"):
if len(messages) > 2:
print("undo latest prompt")
messages = messages[:-2]
else:
print("cannot undo because there are not enough messages on history.")
print("")
if input_text.startswith("!"):
# append_generate_chat("<tool_response>%s</tool_response>" % input_text[1:], role="tool")
append_generate_chat("%s" % input_text[1:], role="tool") # depending on the chat template the tool response tags must or must not be passed. :(
elif input_text.startswith("/regen"):
if len(messages) >= 2:
print("regenerating message (not working)")
messages = messages[:-1]
seed = random.randint(0, 2**32 - 1) # Generate a random seed
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
elif input_text.startswith("/clear"):
print("clearing chat history")
start_msg = messages[0]
messages = [start_msg]
print("")
elif input_text.startswith("/history"):
history = tokenizer.apply_chat_template(messages, return_tensors="pt", tokenize=False, add_generation_prompt=False, tools=tool_list)
print(history)
elif input_text.startswith("/undo"):
if len(messages) > 2:
print("undo latest prompt")
messages = messages[:-2]
else:
print("cannot undo because there are not enough messages on history.")
print("")
elif input_text.startswith("/regen"):
if len(messages) >= 2:
print("regenerating message (not working)")
messages = messages[:-1]
seed = random.randint(0, 2**32 - 1) # Generate a random seed
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
append_generate_chat(None)
else:
print("cannot regenerate because there are not enough messages on history.")
print("")
elif input_text.startswith("/more"):
append_generate_chat(None)
elif input_text.startswith("/file"):
filename = input_text[len("/file "):]
print("read '%s' for prompt:" % filename)
with open(filename, "r") as f:
content = f.read()
print(content)
append_generate_chat(content)
elif input_text.startswith("/auto"):
messages_backup = messages
messages = [roleflip]
for m in messages_backup:
role = m["role"]
content = m["content"]
if role == "user":
role = "assistant"
elif role == "assistant":
role = "user"
if role != "system":
messages.append({"role": role, "content": content})
append_generate_chat(None) # will automatically advance the conversation as 'user'
last_message = messages[-1]
last_message["role"] = "user"
messages = messages_backup + [last_message]
append_generate_chat(None) # 'regular' chatbot answer
elif input_text.startswith("/help"):
print("!<prompt> answer as 'tool' in <tool_response> tags")
print("/clear clear chat history")
print("/undo undo latest prompt")
print("/regen regenerate the last message")
print("/more generate more additional information")
print("/file read prompt input from file")
print("/auto automatically advance conversation")
print("/help print this message")
print("")
elif input_text.startswith("/"):
print("unknown command.")
else:
print("cannot regenerate because there are not enough messages on history.")
print("")
append_generate_chat(input_text)
elif input_text.startswith("/more"):
append_generate_chat(None)
elif input_text.startswith("/auto"):
messages_backup = messages
messages = [roleflip]
for m in messages_backup:
role = m["role"]
content = m["content"]
if role == "user":
role = "assistant"
elif role == "assistant":
role = "user"
if role != "system":
messages.append({"role": role, "content": content})
append_generate_chat(None) # will automatically advance the conversation as 'user'
last_message = messages[-1]
last_message["role"] = "user"
messages = messages_backup + [last_message]
append_generate_chat(None) # 'regular' chatbot answer
elif input_text.startswith("/help"):
print("!<prompt> answer as 'tool' in <tool_response> tags")
print("/clear clear chat history")
print("/undo undo latest prompt")
print("/regen regenerate the last message")
print("/more generate more additional information")
print("/auto automatically advance conversation")
print("/help print this message")
print("")
elif input_text.startswith("/"):
print("unknown command.")
else:
append_generate_chat(input_text)

1
tests/__init__.py Normal file
View File

@@ -0,0 +1 @@
# empty

7
tests/helper.py Normal file
View File

@@ -0,0 +1,7 @@
def tool_dummy(a: int, b: str):
return "result_%d_%s" % (a, b)
def tool_dummy2(text: str):
return text.upper()

30
tests/test_tool_function_decorator.py Normal file
View File

@@ -0,0 +1,30 @@
import pytest
import tool_helper
import tests.helper as helper
def test_tool_function_decorator_if_clean_tool_list():
""" tests for the tool list to be empty. NOT strictly nessesary,
but I want to be warned if this is not the case anymore. Could be not the intention """
start_len = len(tool_helper.tool_list)
assert start_len == 0
def test_tool_function_decorator():
# get length before adding tools
start_len = len(tool_helper.tool_list)
# add tools like it would be a decorator
tool_helper.tool(helper.tool_dummy)
tool_helper.tool(helper.tool_dummy2)
# get length after adding tools
end_len = len(tool_helper.tool_list)
# remove the added ones again
tool_helper.tool_list = tool_helper.tool_list[:-2]
assert end_len == start_len + 2
assert len(tool_helper.tool_list) == start_len

89
tests/test_tool_parse_exec.py Normal file
View File

@@ -0,0 +1,89 @@
import pytest
import tool_helper
from unittest import mock
import tests.helper as helper
def test_tool_dummy():
with mock.patch("tests.helper.tool_dummy") as mock_dummy:
helper.tool_dummy()
mock_dummy.assert_called_once() # this will check if the mocked function on the context was called.
def test_tool_parse_no_exec():
with mock.patch("tests.helper.tool_dummy") as mock_dummy:
tool_helper.parse_and_execute_tool_call("something else", [helper.tool_dummy, helper.tool_dummy2])
assert mock_dummy.call_count == 0
def test_match_and_extract_no_match():
result = tool_helper._match_and_extract("something else", r"<tool_call>(.*)<\/tool_call>")
assert result == None
def test_match_and_extract_matching():
result = tool_helper._match_and_extract("asdfsdfas <tool_call>{json content}</tool_call> adfafsd", r"<tool_call>(.*)<\/tool_call>")
assert result == "{json content}"
def test_match_and_extract_matching2():
result = tool_helper._match_and_extract("<tool_call>{json content}</tool_call>", r"<tool_call>(.*)<\/tool_call>")
assert result == "{json content}"
def test_match_and_extract_matching3_with_newline():
result = tool_helper._match_and_extract("<tool_call>\n{json content}\n</tool_call>", r"<tool_call>(.*)<\/tool_call>")
assert result == "\n{json content}\n"
def test_string_malformed_faulty():
with mock.patch("utils.print_error") as print_error_mock:
result = tool_helper._execute_tool_call_str("{json_content}", [])
assert result == None
print_error_mock.assert_called_once() # this will check if the mocked function on the context was called.
def test_tool_call_json_1():
with mock.patch("utils.print_error") as print_error_mock:
result = tool_helper._execute_tool_call_json({"name": "tool_dummy", "arguments": {"a": 1, "b": "zwei"}}, [helper.tool_dummy, helper.tool_dummy2])
assert result == "result_1_zwei"
assert print_error_mock.call_count == 0
def test_tool_call_json_2():
with mock.patch("utils.print_error") as print_error_mock:
result = tool_helper._execute_tool_call_json({"name": "tool_dummy2", "arguments": {"text": "some_text"}}, [helper.tool_dummy, helper.tool_dummy2])
assert result == "SOME_TEXT"
assert print_error_mock.call_count == 0
def test_tool_call_json_non_existing_call_check():
with mock.patch("utils.print_error") as print_error_mock:
result = tool_helper._execute_tool_call_json({"name": "tool_dummy_which_is_not_existing", "arguments": {"text": "some_text"}}, [helper.tool_dummy, helper.tool_dummy2])
assert result == None
assert print_error_mock.call_count == 1 # this will check if the mocked function on the context was called.
def test_tool_call_json_wrong_arguments_check():
with mock.patch("utils.print_error") as print_error_mock:
result = tool_helper._execute_tool_call_json({"name": "tool_dummy", "arguments": {"a": "must_be_an_int_but_is_string", "b": "zwei"}}, [helper.tool_dummy, helper.tool_dummy2])
assert result == None
assert print_error_mock.call_count == 1 # this will check if the mocked function on the context was called.
def test_regex_multiline():
import re
pattern = r"<start>(.*)</end>"
# The text to search (spanning multiple lines)
text = """<start>
{json}
</end>"""
# Use re.search with re.DOTALL to match across newlines
match = re.search(pattern, text, re.DOTALL)
assert match.group(1).find("{json}") != -1

35
tool_functions.py Normal file
View File

@@ -0,0 +1,35 @@
import random
import datetime
from tool_helper import tool
@tool
def current_time():
"""Get the current local date and time as a string."""
return datetime.datetime.now().strftime("%Y-%m-%d %H:%M")
# @tool
# def random_float():
# """Generate a random float from 0..1."""
# return str(random.random())
@tool
def random_float(a: float=0.0, b: float=1.0):
"""Generate a random float in range [a, b], including both end points. Optional pass no parameter and range 0..1 will be used.
Args:
a: minimum possible value
b: maximum possible value"""
return str(random.randint(a, b))
@tool
def random_int(a: int, b: int):
"""Generate a random integer in range [a, b], including both end points.
Args:
a: minimum possible value
b: maximum possible value"""
return str(random.randint(a, b))
def register_dummy():
pass # dummy function to run and be sure the decorators have run

102
tool_helper.py Normal file
View File

@@ -0,0 +1,102 @@
from typing import Callable, List, Optional
import json
import re
import utils
tool_list = []
def tool(fn):
"""tool function decorator"""
print("register tool '%s'" % fn.__name__)
tool_list.append(fn)
# def parse_and_execute_tool_call(message: str, tools: list[function]) -> str | None:
# """execute tool call if needed accordint <tool_call> tag and return the content of the tool call or None if no call happened."""
# pass
def parse_and_execute_tool_call(message: str, tools: List[Callable]) -> Optional[str]:
"""
Execute a tool call if the <tool_call> tag is present and return the tool's response.
If no <tool_call> tag is found, return None.
Args:
message (str): The message containing the tool call.
tools (list[function]): A list of tool functions available for execution.
Returns:
Optional[str]: The content of the tool response or None if no tool call occurred.
"""
# in case LLM responds with <tool_call></tool_call> the correct way
extracted = _match_and_extract(message, r"<tool_call>(.*)<\/tool_call>")
if extracted:
return _execute_tool_call_str(extracted, tools)
# in case LLM responds with <tool_call></tool_response> by accident
extracted = _match_and_extract(message, r"<tool_call>(.*)<\/tool_.*>")
if extracted:
return _execute_tool_call_str(extracted, tools)
# in case LLM responds with <tool_call></???> by accident
extracted = _match_and_extract(message, r"<tool_call>(.*)<\/.*>")
if extracted:
return _execute_tool_call_str(extracted, tools)
# in case LLM responds with <tool_call></???> by accident
extracted = _match_and_extract(message, r"<tool_response>(.*)<\/.*>")
if extracted:
return _execute_tool_call_str(extracted, tools)
return None
def _match_and_extract(message: str, pattern: str) -> Optional[str]:
""" helper function to match regex and extract group 1 """
match = re.search(pattern, message, re.DOTALL)
if match:
group1 = match.group(1)
return group1
return None
def _execute_tool_call_str(tool_call: str, tools: List[Callable]) -> Optional[str]:
""" execute tool call per string content. The content must be a valid json """
try:
js = json.loads(tool_call)
return _execute_tool_call_json(js, tools)
except json.JSONDecodeError:
utils.print_error("Json was malformed. Will be ignored.")
return None
def _execute_tool_call_json(data: any, tools: List[Callable]) -> Optional[str]:
""" extract name and arguments from parsed data and call the tool, which is matched from the tools list """
# Extract tool name and arguments
tool_name = data.get("name")
arguments = data.get("arguments", {})
# Find the tool by name in the list of tools
for tool in tools:
if tool.__name__ == tool_name:
# Execute the tool
return _execute_tool_function(arguments, tool)
utils.print_error("Specified tool '%s' not found." % tool_name)
return None
def _execute_tool_function(arguments: any, tool: Callable) -> Optional[str]:
""" Execute the tool and return the result. """
try:
result = tool(**arguments)
print("<tool_response>", result, "</tool_response>")
return result
except TypeError as e:
utils.print_error("Type error while executing function call: '%s'" % str(e))
return None

11
utils.py Normal file
View File

@@ -0,0 +1,11 @@
import json
import sys
def load_json_file(filepath: str) -> any:
with open(filepath, "r") as f:
return json.load(f)
def print_error(*args, **kwargs):
print(*args, file=sys.stderr, **kwargs)