Add documentation for updated chat template

README.md

@@ -21,13 +21,37 @@ should probably proofread and complete it, then remove this comment. -->
 
 This model is a fine-tuned version of [Qwen/Qwen2-1.5B-Instruct](https://huggingface.co/Qwen/Qwen2-1.5B-Instruct) on [devanshamin/gem-viggo-function-calling](https://huggingface.co/datasets/devanshamin/gem-viggo-function-calling) dataset.
 
+## Updated Chat Template
+> Note: The template supports multiple tools but the model is fine-tuned on a dataset consisting of a single tool.
+
+- The chat template has been added to the [tokenizer_config.json](https://huggingface.co/devanshamin/Qwen2-1.5B-Instruct-Function-Calling-v1/blob/7ee7c020cefdb0101939469de608acc2afa7809e/tokenizer_config.json#L34).
+- Supports prompts with and without tools.
+
+```python
+chat_template = (
+  "{% for message in messages %}"
+  "{% if loop.first and messages[0]['role'] != 'system' %}"
+  "{% if tools %}"
+  "<|im_start|>system\nYou are a helpful assistant with access to the following tools. Use them if required - \n"
+  "```json\n{{ tools | tojson }}\n```<|im_end|>\n"
+  "{% else %}"
+  "<|im_start|>system\nYou are a helpful assistant.\n<|im_end|>\n"
+  "{% endif %}"
+  "{% endif %}"
+  "{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}"
+  "{% endfor %}"
+  "{% if add_generation_prompt %}"
+  "{{ '<|im_start|>assistant\n' }}"
+  "{% endif %}"
+)
+```
+
 ## Basic Usage
 
 ```python
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-# Load the model and the tokenizer
 model_id = "Qwen2-1.5B-Instruct-Function-Calling-v1"
 model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float32, device_map="auto")
 tokenizer = AutoTokenizer.from_pretrained(model_id)
@@ -39,11 +63,7 @@ def inference(prompt: str) -> str:
   response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
   return response
 
-prompt = "What is the meaning of life?"
-messages = [
-  {"role": "system", "content": "You are a helpful assistant."},
-  {"role": "user", "content": prompt}
-]
+messages = [{"role": "user", "content": "What is the speed of light?"}]
 prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
 response = inference(prompt)
 print(response)
@@ -55,14 +75,17 @@ print(response)
 
 ```python
 import json
-
-def get_prompt(tool: str, user_input: str) -> str:
-  system = "You are a helpful assistant with access to the following tools. Use them if required - \n```json\n{}\n```"
-  messages = [
-    {"role": "system", "content": system.format(tool)},
-    {"role": "user", "content": 'Extract the information from the following - \n{}'.format(user_input)}
-  ]
-  prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+from typing import List, Dict
+
+def get_prompt(user_input: str, tools: List[Dict] | None = None):
+  prompt = 'Extract the information from the following - \n{}'.format(user_input)
+  messages = [{"role": "user", "content": prompt}]
+  prompt = tokenizer.apply_chat_template(
+    messages,
+    tokenize=False,
+    add_generation_prompt=True,
+    tools=tools
+  )
   return prompt
 
 tool = {
@@ -87,7 +110,7 @@ tool = {
   }
 }
 input_text = "Founded in 2021, Pluto raised $4 million across multiple seed funding rounds, valuing the company at $12 million (pre-money), according to PitchBook. The startup was backed by investors including Switch Ventures, Caffeinated Capital and Maxime Seguineau."
-prompt = get_prompt(json.dumps(tool), input_text)
+prompt = get_prompt(input_text, tools=[tool])
 response = inference(prompt)
 print(response)
 # ```json
@@ -100,7 +123,7 @@ print(response)
 #       "Caffeinated Capital",
 #       "Maxime Seguineau"
 #     ],
-#     "valuation": "pre-money $12M",
+#     "valuation": "$12 million",
 #     "source": "PitchBook"
 #   }
 # }
@@ -127,7 +150,7 @@ class Classification(BaseModel):
 function_definition = openai_schema(Classification).openai_schema
 tool = dict(type='function', function=function_definition)
 input_text = "1,25-dihydroxyvitamin D(3) (1,25(OH)(2)D(3)), the biologically active form of vitamin D, is widely recognized as a modulator of the immune system as well as a regulator of mineral metabolism. The objective of this study was to determine the effects of vitamin D status and treatment with 1,25(OH)(2)D(3) on diabetes onset in non-obese diabetic (NOD) mice, a murine model of human type I diabetes. We have found that vitamin D-deficiency increases the incidence of diabetes in female mice from 46% (n=13) to 88% (n=8) and from 0% (n=10) to 44% (n=9) in male mice as of 200 days of age when compared to vitamin D-sufficient animals. Addition of 50 ng of 1,25(OH)(2)D(3)/day to the diet prevented disease onset as of 200 days and caused a significant rise in serum calcium levels, regardless of gender or vitamin D status. Our results indicate that vitamin D status is a determining factor of disease susceptibility and oral administration of 1,25(OH)(2)D(3) prevents diabetes onset in NOD mice through 200 days of age."
-prompt = get_prompt(json.dumps(tool), input_text)
+prompt = get_prompt(input_text, tools=[tool])
 output = inference(prompt)
 print(output)
 # ```json