Skip to content

Commit 9f9afbb

Browse files
hwchase17hwchase17
authored
add custom prompt for LLMMathChain and SQLDatabase chain (#605)
1 parent a87a2aa commit 9f9afbb

File tree

4 files changed

+192
-8
lines changed

4 files changed

+192
-8
lines changed

‎docs/modules/chains/examples/llm_math.ipynb

Lines changed: 99 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212
},
1313
{
1414
"cell_type": "code",
15-
"execution_count": 5,
15+
"execution_count": 1,
1616
"id": "44e9ba31",
1717
"metadata": {},
1818
"outputs": [
@@ -31,7 +31,7 @@
3131
"\u001b[0m\n",
3232
"Answer: \u001b[33;1m\u001b[1;3m2.4116004626599237\n",
3333
"\u001b[0m\n",
34-
"\u001b[1m> Finished LLMMathChain chain.\u001b[0m\n"
34+
"\u001b[1m> Finished chain.\u001b[0m\n"
3535
]
3636
},
3737
{
@@ -40,7 +40,7 @@
4040
"'Answer: 2.4116004626599237\\n'"
4141
]
4242
},
43-
"execution_count": 5,
43+
"execution_count": 1,
4444
"metadata": {},
4545
"output_type": "execute_result"
4646
}
@@ -54,10 +54,105 @@
5454
"llm_math.run(\"What is 13 raised to the .3432 power?\")"
5555
]
5656
},
57+
{
58+
"cell_type": "markdown",
59+
"id": "2bdd5fc6",
60+
"metadata": {},
61+
"source": [
62+
"## Customize Prompt\n",
63+
"You can also customize the prompt that is used. Here is an example prompting it to use numpy"
64+
]
65+
},
66+
{
67+
"cell_type": "code",
68+
"execution_count": 24,
69+
"id": "76be17b0",
70+
"metadata": {},
71+
"outputs": [],
72+
"source": [
73+
"from langchain.prompts.prompt import PromptTemplate\n",
74+
"\n",
75+
"_PROMPT_TEMPLATE = \"\"\"You are GPT-3, and you can't do math.\n",
76+
"\n",
77+
"You can do basic math, and your memorization abilities are impressive, but you can't do any complex calculations that a human could not do in their head. You also have an annoying tendency to just make up highly specific, but wrong, answers.\n",
78+
"\n",
79+
"So we hooked you up to a Python 3 kernel, and now you can execute code. If you execute code, you must print out the final answer using the print function. You MUST use the python package numpy to answer your question. You must import numpy as np.\n",
80+
"\n",
81+
"\n",
82+
"Question: ${{Question with hard calculation.}}\n",
83+
"```python\n",
84+
"${{Code that prints what you need to know}}\n",
85+
"print(${{code}})\n",
86+
"```\n",
87+
"```output\n",
88+
"${{Output of your code}}\n",
89+
"```\n",
90+
"Answer: ${{Answer}}\n",
91+
"\n",
92+
"Begin.\n",
93+
"\n",
94+
"Question: What is 37593 * 67?\n",
95+
"\n",
96+
"```python\n",
97+
"import numpy as np\n",
98+
"print(np.multiply(37593, 67))\n",
99+
"```\n",
100+
"```output\n",
101+
"2518731\n",
102+
"```\n",
103+
"Answer: 2518731\n",
104+
"\n",
105+
"Question: {question}\"\"\"\n",
106+
"\n",
107+
"PROMPT = PromptTemplate(input_variables=[\"question\"], template=_PROMPT_TEMPLATE)"
108+
]
109+
},
110+
{
111+
"cell_type": "code",
112+
"execution_count": 25,
113+
"id": "0c42faa0",
114+
"metadata": {},
115+
"outputs": [
116+
{
117+
"name": "stdout",
118+
"output_type": "stream",
119+
"text": [
120+
"\n",
121+
"\n",
122+
"\u001b[1m> Entering new LLMMathChain chain...\u001b[0m\n",
123+
"What is 13 raised to the .3432 power?\u001b[32;1m\u001b[1;3m\n",
124+
"\n",
125+
"```python\n",
126+
"import numpy as np\n",
127+
"print(np.power(13, .3432))\n",
128+
"```\n",
129+
"\u001b[0m\n",
130+
"Answer: \u001b[33;1m\u001b[1;3m2.4116004626599237\n",
131+
"\u001b[0m\n",
132+
"\u001b[1m> Finished chain.\u001b[0m\n"
133+
]
134+
},
135+
{
136+
"data": {
137+
"text/plain": [
138+
"'Answer: 2.4116004626599237\\n'"
139+
]
140+
},
141+
"execution_count": 25,
142+
"metadata": {},
143+
"output_type": "execute_result"
144+
}
145+
],
146+
"source": [
147+
"llm_math = LLMMathChain(llm=llm, prompt=PROMPT, verbose=True)\n",
148+
"\n",
149+
"llm_math.run(\"What is 13 raised to the .3432 power?\")"
150+
]
151+
},
57152
{
58153
"cell_type": "code",
59154
"execution_count": null,
60-
"id": "f62f0c75",
155+
"id": "0c62951b",
61156
"metadata": {},
62157
"outputs": [],
63158
"source": []

‎docs/modules/chains/examples/sqlite.ipynb

Lines changed: 85 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -78,7 +78,7 @@
7878
"SQLQuery:\u001b[32;1m\u001b[1;3m SELECT COUNT(*) FROM Employee;\u001b[0m\n",
7979
"SQLResult: \u001b[33;1m\u001b[1;3m[(9,)]\u001b[0m\n",
8080
"Answer:\u001b[32;1m\u001b[1;3m There are 9 employees.\u001b[0m\n",
81-
"\u001b[1m> Finished SQLDatabaseChain chain.\u001b[0m\n"
81+
"\u001b[1m> Finished chain.\u001b[0m\n"
8282
]
8383
},
8484
{
@@ -96,10 +96,93 @@
9696
"db_chain.run(\"How many employees are there?\")"
9797
]
9898
},
99+
{
100+
"cell_type": "markdown",
101+
"id": "aad2cba6",
102+
"metadata": {},
103+
"source": [
104+
"## Customize Prompt\n",
105+
"You can also customize the prompt that is used. Here is an example prompting it to understand that foobar is the same as the Employee table"
106+
]
107+
},
108+
{
109+
"cell_type": "code",
110+
"execution_count": 5,
111+
"id": "8ca7bafb",
112+
"metadata": {},
113+
"outputs": [],
114+
"source": [
115+
"from langchain.prompts.prompt import PromptTemplate\n",
116+
"\n",
117+
"_DEFAULT_TEMPLATE = \"\"\"Given an input question, first create a syntactically correct {dialect} query to run, then look at the results of the query and return the answer.\n",
118+
"Use the following format:\n",
119+
"\n",
120+
"Question: \"Question here\"\n",
121+
"SQLQuery: \"SQL Query to run\"\n",
122+
"SQLResult: \"Result of the SQLQuery\"\n",
123+
"Answer: \"Final answer here\"\n",
124+
"\n",
125+
"Only use the following tables:\n",
126+
"\n",
127+
"{table_info}\n",
128+
"\n",
129+
"If someone asks for the table foobar, they really mean the employee table.\n",
130+
"\n",
131+
"Question: {input}\"\"\"\n",
132+
"PROMPT = PromptTemplate(\n",
133+
" input_variables=[\"input\", \"table_info\", \"dialect\"], template=_DEFAULT_TEMPLATE\n",
134+
")"
135+
]
136+
},
137+
{
138+
"cell_type": "code",
139+
"execution_count": 6,
140+
"id": "ec47a2bf",
141+
"metadata": {},
142+
"outputs": [],
143+
"source": [
144+
"db_chain = SQLDatabaseChain(llm=llm, database=db, prompt=PROMPT, verbose=True)"
145+
]
146+
},
147+
{
148+
"cell_type": "code",
149+
"execution_count": 7,
150+
"id": "ebb0674e",
151+
"metadata": {},
152+
"outputs": [
153+
{
154+
"name": "stdout",
155+
"output_type": "stream",
156+
"text": [
157+
"\n",
158+
"\n",
159+
"\u001b[1m> Entering new SQLDatabaseChain chain...\u001b[0m\n",
160+
"How many employees are there in the foobar table? \n",
161+
"SQLQuery:\u001b[32;1m\u001b[1;3m SELECT COUNT(*) FROM Employee;\u001b[0m\n",
162+
"SQLResult: \u001b[33;1m\u001b[1;3m[(9,)]\u001b[0m\n",
163+
"Answer:\u001b[32;1m\u001b[1;3m There are 9 employees in the foobar table.\u001b[0m\n",
164+
"\u001b[1m> Finished chain.\u001b[0m\n"
165+
]
166+
},
167+
{
168+
"data": {
169+
"text/plain": [
170+
"' There are 9 employees in the foobar table.'"
171+
]
172+
},
173+
"execution_count": 7,
174+
"metadata": {},
175+
"output_type": "execute_result"
176+
}
177+
],
178+
"source": [
179+
"db_chain.run(\"How many employees are there in the foobar table?\")"
180+
]
181+
},
99182
{
100183
"cell_type": "code",
101184
"execution_count": null,
102-
"id": "61d91b85",
185+
"id": "e59a4740",
103186
"metadata": {},
104187
"outputs": [],
105188
"source": []

‎langchain/chains/llm_math/base.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
from langchain.chains.llm import LLMChain
88
from langchain.chains.llm_math.prompt import PROMPT
99
from langchain.llms.base import BaseLLM
10+
from langchain.prompts.base import BasePromptTemplate
1011
from langchain.python import PythonREPL
1112

1213

@@ -22,6 +23,8 @@ class LLMMathChain(Chain, BaseModel):
2223

2324
llm: BaseLLM
2425
"""LLM wrapper to use."""
26+
prompt: BasePromptTemplate = PROMPT
27+
"""Prompt to use to translate to python if neccessary."""
2528
input_key: str = "question" #: :meta private:
2629
output_key: str = "answer" #: :meta private:
2730

@@ -48,7 +51,7 @@ def output_keys(self) -> List[str]:
4851
return [self.output_key]
4952

5053
def _call(self, inputs: Dict[str, str]) -> Dict[str, str]:
51-
llm_executor = LLMChain(prompt=PROMPT, llm=self.llm)
54+
llm_executor = LLMChain(prompt=self.prompt, llm=self.llm)
5255
python_executor = PythonREPL()
5356
if self.verbose:
5457
self.callback_manager.on_text(inputs[self.input_key])

‎langchain/chains/sql_database/base.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
from langchain.chains.llm import LLMChain
88
from langchain.chains.sql_database.prompt import PROMPT
99
from langchain.llms.base import BaseLLM
10+
from langchain.prompts.base import BasePromptTemplate
1011
from langchain.sql_database import SQLDatabase
1112

1213

@@ -25,6 +26,8 @@ class SQLDatabaseChain(Chain, BaseModel):
2526
"""LLM wrapper to use."""
2627
database: SQLDatabase
2728
"""SQL Database to connect to."""
29+
prompt: BasePromptTemplate = PROMPT
30+
"""Prompt to use to translate natural language to SQL."""
2831
input_key: str = "query" #: :meta private:
2932
output_key: str = "result" #: :meta private:
3033

@@ -51,7 +54,7 @@ def output_keys(self) -> List[str]:
5154
return [self.output_key]
5255

5356
def _call(self, inputs: Dict[str, str]) -> Dict[str, str]:
54-
llm_chain = LLMChain(llm=self.llm, prompt=PROMPT)
57+
llm_chain = LLMChain(llm=self.llm, prompt=self.prompt)
5558
input_text = f"{inputs[self.input_key]} \nSQLQuery:"
5659
if self.verbose:
5760
self.callback_manager.on_text(input_text)

0 commit comments

Comments
 (0)