Update modeling_functionary.py

modeling_functionary.py

@@ -66,7 +66,7 @@ class FunctionaryForCausalLM(LlamaForCausalLM):
         )
 
         input_ids = kwargs.pop("input_ids")
-        function_call_token = "<|reserved_special_token_249|>"
+        function_call_token = "<function="
         
         correct_results = []
         for input_id, result in zip(input_ids, results):
@@ -80,13 +80,14 @@ class FunctionaryForCausalLM(LlamaForCausalLM):
                     continue
 
                 chunk = chunk.replace(tokenizer.pad_token, "")
+
                 if i == 0 and has_text is not False:
-                    final_output_json["content"] = chunk.strip[:-len("<|eot_id|>")] if chunk.endswith("<|eot_id|>") else chunk
+                    final_output_json["content"] = chunk.removesuffix("<|eom_id|>").removesuffix("<|eot_id|>")
                 else:
                     tool_calls.append(
                         {
-                            "name": chunk[: chunk.index("\n{")],
-                            "arguments": chunk[chunk.index("\n{") + 1: -len("<|eot_id|>")] if chunk.endswith("<|eot_id|>") else chunk[chunk.index("\n{") + 1:]
+                            "name": chunk[: chunk.index(">{")],
+                            "arguments": chunk[chunk.index(">{") + 1: ].removesuffix("<|eom_id|>").removesuffix("</function>")
                         }
                     )
             if len(tool_calls) > 0: