File size: 7,443 Bytes
59699eb
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
import json
import time
import argparse
from typing import List, Dict

from vllm import LLM, SamplingParams
from jinja2 import Template


# Default prompts
TASK_INSTRUCTION = """
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.
""".strip()

FORMAT_INSTRUCTION = """
The output MUST strictly adhere to the following JSON 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 tool_calls an empty list '[]'
```
{
  "tool_calls": [
    {"name": "func_name1", "arguments": {"argument1": "value1", "argument2": "value2"}},
    ... (more tool calls as required)
  ]
}
```
""".strip()


class XLAMHandler:
    def __init__(self, model: str, temperature: float = 0.3, top_p: float = 1, max_tokens: int = 512):
        self.llm = LLM(model=model)
        self.sampling_params = SamplingParams(
            temperature=temperature,
            top_p=top_p,
            max_tokens=max_tokens
        )
        self.chat_template = self.llm.get_tokenizer().chat_template
    
    @staticmethod
    def apply_chat_template(template, messages):
        jinja_template = Template(template)
        return jinja_template.render(messages=messages)

    def process_query(self, query: str, tools: Dict, task_instruction: str, format_instruction: str):
        # Convert tools to XLAM format
        xlam_tools = self.convert_to_xlam_tool(tools)

        # Build the input prompt
        prompt = self.build_prompt(query, xlam_tools, task_instruction, format_instruction)

        messages = [
            {"role": "user", "content": prompt}
        ]
        formatted_prompt = self.apply_chat_template(self.chat_template, messages)

        # Make inference
        start_time = time.time()
        outputs = self.llm.generate([formatted_prompt], self.sampling_params)
        latency = time.time() - start_time

        # Parse response
        result = outputs[0].outputs[0].text
        parsed_result, success, _ = self.parse_response(result)

        # Prepare metadata
        metadata = {
            "latency": latency,
            "success": success,
        }

        return parsed_result, metadata

    def convert_to_xlam_tool(self, tools):
        if isinstance(tools, dict):
            return {
                "name": tools["name"],
                "description": tools["description"],
                "parameters": {k: v for k, v in tools["parameters"].get("properties", {}).items()}
            }
        elif isinstance(tools, list):
            return [self.convert_to_xlam_tool(tool) for tool in tools]
        else:
            return tools

    def build_prompt(self, query, tools, task_instruction=TASK_INSTRUCTION, format_instruction=FORMAT_INSTRUCTION):
        prompt = f"[BEGIN OF TASK INSTRUCTION]\n{task_instruction}\n[END OF TASK INSTRUCTION]\n\n"
        prompt += f"[BEGIN OF AVAILABLE TOOLS]\n{json.dumps(tools)}\n[END OF AVAILABLE TOOLS]\n\n"
        prompt += f"[BEGIN OF FORMAT INSTRUCTION]\n{format_instruction}\n[END OF FORMAT INSTRUCTION]\n\n"
        prompt += f"[BEGIN OF QUERY]\n{query}\n[END OF QUERY]\n\n"
        return prompt

    def parse_response(self, response):
        try:
            data = json.loads(response)
            tool_calls = data.get('tool_calls', []) if isinstance(data, dict) else data
            result = [
                {tool_call['name']: tool_call['arguments']}
                for tool_call in tool_calls if isinstance(tool_call, dict)
            ]
            return result, True, []
        except json.JSONDecodeError:
            return [], False, ["Failed to parse JSON response"]
        

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Test XLAM model with vLLM")
    parser.add_argument("--model", required=True, help="Path to the model")
    parser.add_argument("--temperature", type=float, default=0.3, help="Temperature for sampling")
    parser.add_argument("--top_p", type=float, default=1.0, help="Top p for sampling")
    parser.add_argument("--max_tokens", type=int, default=512, help="Maximum number of tokens to generate")
    
    args = parser.parse_args()

    # Initialize the XLAMHandler with command-line arguments
    handler = XLAMHandler(args.model, temperature=args.temperature, top_p=args.top_p, max_tokens=args.max_tokens)

    # Test case 1: Weather API, follows the OpenAI format: https://platform.openai.com/docs/guides/function-calling
    weather_api = {
        "name": "get_weather",
        "description": "Get the current weather for a location",
        "parameters": {
            "type": "object",
            "properties": {
                "location": {
                    "type": "string",
                    "description": "The city and state, e.g. San Francisco, CA"
                },
                "unit": {
                    "type": "string",
                    "enum": ["celsius", "fahrenheit"],
                    "description": "The unit of temperature to return"
                }
            },
            "required": ["location"]
        }
    }

    # Test queries
    test_queries = [
        "What's the weather like in New York?",
        "Tell me the temperature in London in Celsius",
        "What's the weather forecast for Tokyo?",
        "What is the stock price of CRM?",  # the model should return an empty list
        "What's the current temperature in Paris in Fahrenheit?"
    ]

    # Run test cases
    for query in test_queries:
        print(f"Query: {query}")
        result, metadata = handler.process_query(query, weather_api, TASK_INSTRUCTION, FORMAT_INSTRUCTION)
        print(f"Result: {json.dumps(result, indent=2)}")
        print(f"Metadata: {json.dumps(metadata, indent=2)}")
        print("-" * 50)

    # Test case 2: Multiple APIs, follows the OpenAI format: https://platform.openai.com/docs/guides/function-calling
    calculator_api = {
        "name": "calculate",
        "description": "Perform a mathematical calculation",
        "parameters": {
            "type": "object",
            "properties": {
                "operation": {
                    "type": "string",
                    "enum": ["add", "subtract", "multiply", "divide"],
                    "description": "The mathematical operation to perform"
                },
                "x": {
                    "type": "number",
                    "description": "The first number"
                },
                "y": {
                    "type": "number",
                    "description": "The second number"
                }
            },
            "required": ["operation", "x", "y"]
        }
    }

    multi_api_query = "What's the weather in Miami and what's 15 multiplied by 7?"
    multi_api_result, multi_api_metadata = handler.process_query(
        multi_api_query, 
        [weather_api, calculator_api], 
        TASK_INSTRUCTION, 
        FORMAT_INSTRUCTION
    )

    print("Multi-API Query Result:")
    print(json.dumps(multi_api_result, indent=2))
    print(f"Metadata: {json.dumps(multi_api_metadata, indent=2)}")