brian-yu-nexusflow commited on
Commit
ef90054
1 Parent(s): 51e4adc

Upload 2 files

Browse files
Files changed (2) hide show
  1. langchain_example.py +147 -0
  2. non_langchain_example.py +142 -0
langchain_example.py ADDED
@@ -0,0 +1,147 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import List, Literal, Union
2
+
3
+ import math
4
+
5
+ from langchain.tools.base import StructuredTool
6
+ from langchain.agents import (
7
+ Tool,
8
+ AgentExecutor,
9
+ LLMSingleActionAgent,
10
+ AgentOutputParser,
11
+ )
12
+ from langchain.schema import AgentAction, AgentFinish, OutputParserException
13
+ from langchain.prompts import StringPromptTemplate
14
+ from langchain.llms import HuggingFaceTextGenInference
15
+ from langchain.chains import LLMChain
16
+
17
+
18
+ ##########################################################
19
+ # Step 1: Define the functions you want to articulate. ###
20
+ ##########################################################
21
+
22
+
23
+ def calculator(
24
+ input_a: float,
25
+ input_b: float,
26
+ operation: Literal["add", "subtract", "multiply", "divide"],
27
+ ):
28
+ """
29
+ Computes a calculation.
30
+
31
+ Args:
32
+ input_a (float) : Required. The first input.
33
+ input_b (float) : Required. The second input.
34
+ operation (string): The operation. Choices include: add to add two numbers, subtract to subtract two numbers, multiply to multiply two numbers, and divide to divide them.
35
+ """
36
+ match operation:
37
+ case "add":
38
+ return input_a + input_b
39
+ case "subtract":
40
+ return input_a - input_b
41
+ case "multiply":
42
+ return input_a * input_b
43
+ case "divide":
44
+ return input_a / input_b
45
+
46
+
47
+ def cylinder_volume(radius, height):
48
+ """
49
+ Calculate the volume of a cylinder.
50
+
51
+ Parameters:
52
+ - radius (float): The radius of the base of the cylinder.
53
+ - height (float): The height of the cylinder.
54
+
55
+ Returns:
56
+ - float: The volume of the cylinder.
57
+ """
58
+ if radius < 0 or height < 0:
59
+ raise ValueError("Radius and height must be non-negative.")
60
+
61
+ volume = math.pi * (radius**2) * height
62
+ return volume
63
+
64
+
65
+ #############################################################
66
+ # Step 2: Let's define some utils for building the prompt ###
67
+ #############################################################
68
+
69
+
70
+ RAVEN_PROMPT = """
71
+ {raven_tools}
72
+ User Query: Question: {input}
73
+
74
+ Please pick a function from the above options that best answers the user query and fill in the appropriate arguments.<human_end>"""
75
+
76
+
77
+ # Set up a prompt template
78
+ class RavenPromptTemplate(StringPromptTemplate):
79
+ # The template to use
80
+ template: str
81
+ # The list of tools available
82
+ tools: List[Tool]
83
+
84
+ def format(self, **kwargs) -> str:
85
+ prompt = "<human>:\n"
86
+ for tool in self.tools:
87
+ func_signature, func_docstring = tool.description.split(" - ", 1)
88
+ prompt += f'\nOPTION:\n<func_start>def {func_signature}<func_end>\n<docstring_start>\n"""\n{func_docstring}\n"""\n<docstring_end>\n'
89
+ kwargs["raven_tools"] = prompt
90
+ return self.template.format(**kwargs).replace("{{", "{").replace("}}", "}")
91
+
92
+
93
+ class RavenOutputParser(AgentOutputParser):
94
+ def parse(self, llm_output: str) -> Union[AgentAction, AgentFinish]:
95
+ # Check if agent should finish
96
+ if "Initial Answer:" in llm_output:
97
+ return AgentFinish(
98
+ return_values={
99
+ "output": llm_output.strip()
100
+ .split("\n")[1]
101
+ .replace("Initial Answer: ", "")
102
+ .strip()
103
+ },
104
+ log=llm_output,
105
+ )
106
+ else:
107
+ raise OutputParserException(f"Could not parse LLM output: `{llm_output}`")
108
+
109
+
110
+ ##################################################
111
+ # Step 3: Build the agent with these utilities ###
112
+ ##################################################
113
+
114
+
115
+ inference_server_url = "<YOUR ENDPOINT URL>"
116
+ assert (
117
+ inference_server_url is not "<YOUR ENDPOINT URL>"
118
+ ), "Please provide your own HF inference endpoint URL!"
119
+
120
+ llm = HuggingFaceTextGenInference(
121
+ inference_server_url=inference_server_url,
122
+ temperature=0.001,
123
+ max_new_tokens=400,
124
+ do_sample=False,
125
+ )
126
+ tools = [
127
+ StructuredTool.from_function(calculator),
128
+ StructuredTool.from_function(cylinder_volume),
129
+ ]
130
+ raven_prompt = RavenPromptTemplate(
131
+ template=RAVEN_PROMPT, tools=tools, input_variables=["input"]
132
+ )
133
+ llm_chain = LLMChain(llm=llm, prompt=raven_prompt)
134
+ output_parser = RavenOutputParser()
135
+ agent = LLMSingleActionAgent(
136
+ llm_chain=llm_chain,
137
+ output_parser=output_parser,
138
+ stop=["\nReflection:"],
139
+ allowed_tools=tools,
140
+ )
141
+ agent_chain = AgentExecutor.from_agent_and_tools(agent=agent, tools=tools, verbose=True)
142
+
143
+ call = agent_chain.run(
144
+ "I have a cake that is about 3 centimenters high and 200 centimeters in radius. How much cake do I have?"
145
+ )
146
+ call = agent_chain.run("What is 1+10?")
147
+ print(exec(call))
non_langchain_example.py ADDED
@@ -0,0 +1,142 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Literal
2
+
3
+ import math
4
+
5
+ import inspect
6
+
7
+ from transformers import pipeline
8
+
9
+
10
+ ##########################################################
11
+ # Step 1: Define the functions you want to articulate. ###
12
+ ##########################################################
13
+
14
+
15
+ def calculator(
16
+ input_a: float,
17
+ input_b: float,
18
+ operation: Literal["add", "subtract", "multiply", "divide"],
19
+ ):
20
+ """
21
+ Computes a calculation.
22
+
23
+ Args:
24
+ input_a (float) : Required. The first input.
25
+ input_b (float) : Required. The second input.
26
+ operation (string): The operation. Choices include: add to add two numbers, subtract to subtract two numbers, multiply to multiply two numbers, and divide to divide them.
27
+ """
28
+ match operation:
29
+ case "add":
30
+ return input_a + input_b
31
+ case "subtract":
32
+ return input_a - input_b
33
+ case "multiply":
34
+ return input_a * input_b
35
+ case "divide":
36
+ return input_a / input_b
37
+
38
+
39
+ def cylinder_volume(radius, height):
40
+ """
41
+ Calculate the volume of a cylinder.
42
+
43
+ Parameters:
44
+ - radius (float): The radius of the base of the cylinder.
45
+ - height (float): The height of the cylinder.
46
+
47
+ Returns:
48
+ - float: The volume of the cylinder.
49
+ """
50
+ if radius < 0 or height < 0:
51
+ raise ValueError("Radius and height must be non-negative.")
52
+
53
+ volume = math.pi * (radius**2) * height
54
+ return volume
55
+
56
+
57
+ #############################################################
58
+ # Step 2: Let's define some utils for building the prompt ###
59
+ #############################################################
60
+
61
+
62
+ def format_functions_for_prompt(*functions):
63
+ formatted_functions = []
64
+ for func in functions:
65
+ source_code = inspect.getsource(func)
66
+ docstring = inspect.getdoc(func)
67
+ formatted_functions.append(
68
+ f"OPTION:\n<func_start>{source_code}<func_end>\n<docstring_start>\n{docstring}\n<docstring_end>"
69
+ )
70
+ return "\n".join(formatted_functions)
71
+
72
+
73
+ ##############################
74
+ # Step 3: Construct Prompt ###
75
+ ##############################
76
+
77
+
78
+ def construct_prompt(user_query: str):
79
+ formatted_prompt = format_functions_for_prompt(calculator, cylinder_volume)
80
+ formatted_prompt += f"\n\nUser Query: Question: {user_query}\n"
81
+
82
+ prompt = (
83
+ "<human>:\n"
84
+ + formatted_prompt
85
+ + "Please pick a function from the above options that best answers the user query and fill in the appropriate arguments.<human_end>"
86
+ )
87
+ return prompt
88
+
89
+
90
+ #######################################
91
+ # Step 4: Execute the function call ###
92
+ #######################################
93
+
94
+
95
+ def execute_function_call(model_output):
96
+ # Ignore everything after "Reflection" since that is not essential.
97
+ function_call = (
98
+ model_output[0]["generated_text"]
99
+ .strip()
100
+ .split("\n")[1]
101
+ .replace("Initial Answer:", "")
102
+ .strip()
103
+ )
104
+
105
+ try:
106
+ return eval(function_call)
107
+ except Exception as e:
108
+ return str(e)
109
+
110
+
111
+ if __name__ == "__main__":
112
+ # Build the model
113
+ text_gen = pipeline(
114
+ "text-generation",
115
+ model="Nexusflow/NexusRaven-13B",
116
+ device="cuda",
117
+ )
118
+
119
+ # Comp[ute a Simple equation
120
+ prompt = construct_prompt("What is 1+10?")
121
+ model_output = text_gen(
122
+ prompt, do_sample=False, max_new_tokens=400, return_full_text=False
123
+ )
124
+ result = execute_function_call(model_output)
125
+
126
+ print("Model Output:", model_output)
127
+ print("Execution Result:", result)
128
+
129
+ prompt = construct_prompt(
130
+ "I have a cake that is about 3 centimenters high and 200 centimeters in diameter. How much cake do I have?"
131
+ )
132
+ model_output = text_gen(
133
+ prompt,
134
+ do_sample=False,
135
+ max_new_tokens=400,
136
+ return_full_text=False,
137
+ stop=["\nReflection:"],
138
+ )
139
+ result = execute_function_call(model_output)
140
+
141
+ print("Model Output:", model_output)
142
+ print("Execution Result:", result)