Update README.md

README.md

@@ -18,12 +18,126 @@ It has been trained using [TRL](https://github.com/huggingface/trl).
 ## Quick start
 
 ```python
-from transformers import pipeline
+import json
+import re
+from typing import Optional
 
-question = "If you had a time machine, but could only go to the past or the future once and never return, which would you choose and why?"
-generator = pipeline("text-generation", model="plaguss/Llama-3.2-1B-Instruct-v2-FC", device="cuda")
-output = generator([{"role": "user", "content": question}], max_new_tokens=128, return_full_text=False)[0]
-print(output["generated_text"])
+from jinja2 import Template
+import torch 
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers.utils import get_json_schema
+
+
+system_prompt = Template("""You are an expert in composing functions. You are given a question and a set of possible functions. 
+Based on the question, you will need to make one or more function/tool calls to achieve the purpose. 
+If none of the functions can be used, point it out and refuse to answer. 
+If the given question lacks the parameters required by the function, also point it out.
+
+You have access to the following tools:
+<tools>{{ tools }}</tools>
+
+The output MUST strictly adhere to the following format, and NO other text MUST be included.
+The example format is as follows. Please make sure the parameter type is correct. If no function call is needed, please make the tool calls an empty list '[]'.
+<tool_call>[
+{"name": "func_name1", "arguments": {"argument1": "value1", "argument2": "value2"}},
+... (more tool calls as required)
+]</tool_call>""")
+
+
+def prepare_messages(
+    query: str,
+    tools: Optional[dict[str, any]] = None,
+    history: Optional[list[dict[str, str]]] = None
+) -> list[dict[str, str]]:
+    """Prepare the system and user messages for the given query and tools.
+    
+    Args:
+        query: The query to be answered.
+        tools: The tools available to the user. Defaults to None, in which case if a
+            list without content will be passed to the model.
+        history: Exchange of messages, including the system_prompt from
+            the first query. Defaults to None, the first message in a conversation.
+    """
+    if tools is None:
+        tools = []
+    if history:
+        messages = history.copy()
+        messages.append({"role": "user", "content": query})
+    else:
+        messages = [
+            {"role": "system", "content": system_prompt.render(tools=json.dumps(tools))},
+            {"role": "user", "content": query}
+        ]
+    return messages
+
+
+def parse_response(text: str) -> str | dict[str, any]:
+    """Parses a response from the model, returning either the
+    parsed list with the tool calls parsed, or the
+    model thought or response if couldn't generate one.
+
+    Args:
+        text: Response from the model.
+    """
+    pattern = r"<tool_call>(.*?)</tool_call>"
+    matches = re.findall(pattern, text, re.DOTALL)
+    if matches:
+        return json.loads(matches[0])
+    return text
+
+
+model_name_llama = "plaguss/Llama-3.2-1B-Instruct-v2-FC"
+model = AutoModelForCausalLM.from_pretrained(model_name_llama, device_map="auto", torch_dtype="auto", trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained(model_name_llama)
+
+from datetime import datetime
+import random
+
+def get_current_time() -> str:
+    """Returns the current time in 24-hour format.
+
+    Returns:
+        str: Current time in HH:MM:SS format.
+    """
+    return datetime.now().strftime("%H:%M:%S")
+
+
+def get_random_number_between(min: int, max: int) -> int:
+    """
+    Gets a random number between min and max.
+
+    Args:
+        min: The minimum number.
+        max: The maximum number.
+
+    Returns:
+        A random number between min and max.
+    """
+    return random.randint(min, max)
+
+
+tools = [get_json_schema(get_random_number_between), get_json_schema(get_current_time)]
+
+toolbox = {"get_random_number_between": get_random_number_between, "get_current_time": get_current_time}
+
+query = "Give me a number between 1 and 300"
+query = "Can you give me the hour?"
+
+messages = prepare_messages(query, tools=tools)
+
+inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to(model.device)
+outputs = model.generate(inputs, max_new_tokens=512, do_sample=False, num_return_sequences=1, eos_token_id=tokenizer.eos_token_id)
+result = tokenizer.decode(outputs[0][len(inputs[0]):], skip_special_tokens=True)
+
+tool_calls = parse_response(result)
+# [{'name': 'get_random_number_between', 'arguments': {'min': 1, 'max': 300}}
+
+# Get tool responses
+tool_responses = [toolbox.get(tc["name"])(*tc["arguments"].values()) for tc in tool_calls]
+# ['07:20:47']
+
+tool_response = get_random_number_between(*tool_calls[0].get("arguments").values())
+# 45
 ```
 
 ## Training procedure