Implement grammar sampling for function and parameters names

.vscode/settings.json

@@ -3,7 +3,7 @@
     "editor.defaultFormatter": "ms-python.black-formatter",
     "editor.formatOnSave": true,
     "editor.codeActionsOnSave": {
-        "source.organizeImports": true
+      "source.organizeImports": "explicit"
     },
   },
   "isort.args":["--profile", "black"],

functionary/inference.py

@@ -1,12 +1,18 @@
 from typing import List, Optional
 
 import torch
-from transformers import (LlamaForCausalLM, LlamaTokenizer, StoppingCriteria,
-                          StoppingCriteriaList)
+from transformers import (
+    LlamaForCausalLM,
+    LlamaTokenizer,
+    StoppingCriteria,
+    StoppingCriteriaList,
+)
 
 from functionary.openai_types import ChatMessage, Function, FunctionCall, Tool
-from functionary.prompt_template import (PromptTemplate,
-                                         get_prompt_template_from_tokenizer)
+from functionary.prompt_template import (
+    PromptTemplate,
+    get_prompt_template_from_tokenizer,
+)
 
 
 class StopWordsCriteria(StoppingCriteria):
@@ -37,9 +43,8 @@ def prepare_messages_for_inference(
     tools: Optional[List[Tool]] = None,
     device="cuda:0",
 ) -> torch.Tensor:
-
     prompt_template = get_prompt_template_from_tokenizer(tokenizer)
-    
+
     dic_messages = [mess.dict() for mess in messages]
     dic_messages.append({"role": "assistant"})
 

functionary/prompt_template/prompt_template_v1.py

@@ -1,4 +1,5 @@
-from typing import Any, Dict, List, Optional, Tuple, Union
+import json
+from typing import Any, Dict, List, Literal, Optional, Tuple, Union
 
 from functionary.prompt_template.base_template import PromptTemplate
 
@@ -10,6 +11,9 @@ class PromptTemplateV1(PromptTemplate):
     end_assistant = "<|END_OF_ASSISTANT|>"
     end_function = "<|END_OF_FUNCTION_RESULT|>"
     end_function_call = "<|END_OF_FUNCTION_CALL|>"
+    version = "v1"
+    # This token splits between function name and parameters
+    fn_param_sep_token = ":\n{"
 
     def get_end_token_from_message(self, message: Dict) -> str:
         """this function is used for getting the end token for each message.
@@ -36,6 +40,26 @@ def get_end_token_from_message(self, message: Dict) -> str:
             else:
                 return self.end_assistant
 
+    def get_start_of_function_call_token(self) -> str:
+        return self.start_function
+
+    def get_stop_token_for_function_parameter(
+        self, stage: Literal["function", "parameter"]
+    ) -> int:
+        if stage == "function":
+            return ":"  # 28747
+        else:
+            return '":'  # 1264
+
+    def initialize_grammar_sampling_gen_state(self) -> Dict:
+        return {
+            "stage": "pre-function",
+            "curr_tokens": [],
+            "curr_text": "",
+            "func_name": "",
+            "param_names": [],
+        }
+
     def get_additional_tokens(self) -> List[str]:
         return [
             self.start_function,
@@ -234,4 +258,3 @@ def get_chat_template_jinja(self) -> str:
         chat_template = chat_template.replace("<br>\n", "")
         chat_template = chat_template.strip()
         return chat_template
-

functionary/prompt_template/prompt_template_v2.py

@@ -1,6 +1,7 @@
+import json
 import random
 import string
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Literal, Optional, Tuple, Union
 
 from functionary.prompt_template.base_template import PromptTemplate
 
@@ -10,6 +11,32 @@ class PromptTemplateV2(PromptTemplate):
     recipient_token = "<|recipient|>"
     content_token = "<|content|>"
     stop_token = "<|stop|>"
+    version = "v2"
+    # This token splits between function name and parameters
+    fn_param_sep_token = "\n<|content|> {"
+
+    def get_start_of_function_call_token(self) -> str:
+        return self.recipient_token
+
+    def get_stop_token_for_function_parameter(
+        self, stage: Literal["function", "parameter"]
+    ) -> int:
+        if stage == "function":
+            return "\n"  # 13
+        else:
+            return '":'  # 1264
+
+    def get_predefined_function_names(self) -> List[str]:
+        return ["all"]
+
+    def initialize_grammar_sampling_gen_state(self) -> Dict:
+        return {
+            "stage": "function",
+            "curr_tokens": [],
+            "curr_text": "",
+            "func_name": "",
+            "param_names": [],
+        }
 
     def get_additional_tokens(self) -> List[str]:
         return [
@@ -264,7 +291,7 @@ def update_response_state_from_delta_text(
                 "func_name": None,  # function_name of the current tool, if the response requires to use tool
                 "response_type": None,  # response_type=text(text response)/function (using tool)
                 "func_index": -1,  # index of the tool in tool_calls
-                "call_id": None,  # call_id of the current tool 
+                "call_id": None,  # call_id of the current tool
                 # skip_until_reach we skip new tokens until we reach certain token. This is used when we hit special tokens
                 "skip_until_reach": self.content_token,  # at first we will skip until reach <|content|>
                 "first_time": True,  # if first_time we return an tempty delta with role=assistant