upload model

encoding/test_encoding_dsv32.py

@@ -0,0 +1,56 @@
+import json
+import copy
+
+from encoding_dsv32 import encode_messages, parse_message_from_completion_text
+
+with open("test_input.json", "r") as f:
+    test_dict = json.load(f)
+    messages = test_dict["messages"]
+    messages[0]["tools"] = test_dict["tools"]
+
+with open("test_output.txt", "r") as f:
+    gold_prompt = f.read().strip()
+
+print(messages)
+print("=" * 60)
+
+encode_config = dict(thinking_mode="thinking", drop_thinking=True, add_default_bos_token=True)
+prompt = encode_messages(messages, **encode_config)
+print(prompt)
+assert prompt == gold_prompt
+print("=" * 60)
+
+tool_call_message = messages[4]
+tool_call_prompt = encode_messages([tool_call_message], context=messages[:4], **encode_config)
+tool_call_message_wo_id = copy.deepcopy(tool_call_message)
+for tool_call in tool_call_message_wo_id["tool_calls"]:
+    tool_call.pop("id")
+parsed_tool_call_message = parse_message_from_completion_text(tool_call_prompt, thinking_mode="thinking")
+parsed_tool_call_message.pop("content")
+assert tool_call_message_wo_id == parsed_tool_call_message
+
+thinking_message = messages[-6]
+thinking_prompt = encode_messages([thinking_message], context=messages[:-6], **encode_config)
+parsed_thinking_message = parse_message_from_completion_text(thinking_prompt, thinking_mode="thinking")
+parsed_thinking_message.pop("tool_calls")
+assert thinking_message == parsed_thinking_message
+
+with open("test_input_search_wo_date.json", "r") as f:
+    search_messages = json.load(f)["messages"]
+
+with open("test_output_search_wo_date.txt", "r") as f:
+    search_gold_prompt = f.read().strip()
+
+search_prompt = encode_messages(search_messages, **encode_config)
+assert search_prompt == search_gold_prompt
+
+with open("test_input_search_w_date.json", "r") as f:
+    search_messages_w_date = json.load(f)["messages"]
+
+with open("test_output_search_w_date.txt", "r") as f:
+    search_gold_prompt_w_date = f.read().strip()
+
+search_prompt_w_date = encode_messages(search_messages_w_date, **encode_config)
+with open("test_output_search_w_date_2.txt", "w") as f:
+    f.write(search_prompt_w_date)
+assert search_prompt_w_date == search_gold_prompt_w_date