324 lines
11 KiB
Python
324 lines
11 KiB
Python
"""
|
||
Class which manages a Pantograph instance. All calls to the kernel uses this
|
||
interface.
|
||
"""
|
||
import json, pexpect, pathlib, unittest
|
||
from pantograph.expr import Variable, Goal, GoalState, \
|
||
Tactic, TacticHave, TacticCalc
|
||
|
||
|
||
import argparse
|
||
from typing import Dict, List
|
||
import os
|
||
import sglang as sgl
|
||
|
||
|
||
|
||
|
||
@sgl.function
|
||
def multi_turn_question(s, question_1, question_2):
|
||
s += sgl.system("You are a helpful assistant.")
|
||
s += sgl.user(question_1)
|
||
s += sgl.assistant(sgl.gen("answer_1", max_tokens=256))
|
||
s += sgl.user(question_2)
|
||
s += sgl.assistant(sgl.gen("answer_2", max_tokens=256))
|
||
|
||
|
||
@sgl.function
|
||
def select_tactic(s, state):
|
||
s += sgl.system("You are an expert in Lean. Choose the next one tactic to run given the current proof state and goals.")
|
||
s += sgl.user("The current proof state: GoalState(state_id=0, goals=[Goal(variables=[], target='∀ (a b: Nat), (b = 2) -> 1 + a + 1 = a + b', name=None, is_conversion=False)])")
|
||
s += sgl.assistant("```intros a b h```")
|
||
s += sgl.user("The current proof state: GoalState(state_id=1, goals=[Goal(variables=[Variable(t='Nat', v=None, name='a'), Variable(t='Nat', v=None, name='b'), Variable(t='b = 2', v=None, name='h')], target='1 + a + 1 = a + b', name=None, is_conversion=False)])")
|
||
s += sgl.assistant('TacticCalc("1 + a + 1 = a + 1 + 1")')
|
||
s += sgl.user("The current proof state: " + str(state))
|
||
with s.copy() as tmp:
|
||
tmp += sgl.assistant(sgl.gen("tactic", max_tokens=64))
|
||
print("==tmp===")
|
||
print(tmp["tactic"])
|
||
tactic = extract_code_from_llm_output(tmp["tactic"])
|
||
s += sgl.assistant("```"+tactic+"```")
|
||
|
||
|
||
|
||
def _get_proc_cwd():
|
||
return pathlib.Path(__file__).parent
|
||
def _get_proc_path():
|
||
return _get_proc_cwd() / "pantograph"
|
||
|
||
class ServerError(Exception):
|
||
pass
|
||
|
||
class Server:
|
||
|
||
def __init__(self,
|
||
imports=["Init"],
|
||
options=[],
|
||
timeout=20,
|
||
maxread=1000000):
|
||
"""
|
||
timeout: Amount of time to wait for execution
|
||
maxread: Maximum number of characters to read (especially important for large proofs and catalogs)
|
||
"""
|
||
self.timeout = timeout
|
||
self.imports = imports
|
||
self.maxread = maxread
|
||
self.proc_cwd = _get_proc_cwd()
|
||
self.proc_path = _get_proc_path()
|
||
|
||
self.options = options
|
||
self.args = " ".join(imports + [f'--{opt}' for opt in options])
|
||
self.proc = None
|
||
self.restart()
|
||
|
||
def restart(self):
|
||
if self.proc is not None:
|
||
self.proc.close()
|
||
self.proc = pexpect.spawn(
|
||
f"{self.proc_path} {self.args}",
|
||
encoding="utf-8",
|
||
maxread=self.maxread,
|
||
cwd=self.proc_cwd,
|
||
)
|
||
self.proc.setecho(False)
|
||
|
||
def run(self, cmd, payload):
|
||
s = json.dumps(payload)
|
||
self.proc.sendline(f"{cmd} {s}")
|
||
try:
|
||
self.proc.expect("{.*}\r\n", timeout=self.timeout)
|
||
output = self.proc.match.group()
|
||
return json.loads(output)
|
||
except pexpect.exceptions.TIMEOUT:
|
||
raise pexpect.exceptions.TIMEOUT
|
||
|
||
def reset(self):
|
||
return self.run("reset", {})
|
||
|
||
def goal_start(self, expr: str) -> GoalState:
|
||
result = self.run('goal.start', {"expr": str(expr)})
|
||
if "error" in result:
|
||
raise ServerError(result["desc"])
|
||
return GoalState(state_id = result["stateId"], goals = [Goal.sentence(expr)])
|
||
|
||
def goal_tactic(self, state: GoalState, goal_id: int, tactic: Tactic) -> GoalState:
|
||
args = {"stateId": state.state_id, "goalId": goal_id}
|
||
if isinstance(tactic, str):
|
||
args["tactic"] = tactic
|
||
elif isinstance(tactic, TacticHave):
|
||
args["have"] = tactic.branch
|
||
elif isinstance(tactic, TacticCalc):
|
||
args["calc"] = tactic.step
|
||
else:
|
||
raise RuntimeError(f"Invalid tactic type: {tactic}")
|
||
result = self.run('goal.tactic', args)
|
||
if "error" in result:
|
||
raise ServerError(result["desc"])
|
||
if "tacticErrors" in result:
|
||
raise ServerError(result["tacticErrors"])
|
||
if "parseError" in result:
|
||
raise ServerError(result["parseError"])
|
||
return GoalState.parse(result)
|
||
|
||
def goal_conv_begin(self, state: GoalState, goal_id: int) -> GoalState:
|
||
result = self.run('goal.tactic', {"stateId": state.state_id, "goalId": goal_id, "conv": True})
|
||
if "error" in result:
|
||
raise ServerError(result["desc"])
|
||
if "tacticErrors" in result:
|
||
raise ServerError(result["tacticErrors"])
|
||
if "parseError" in result:
|
||
raise ServerError(result["parseError"])
|
||
return GoalState.parse(result)
|
||
|
||
def goal_conv_end(self, state: GoalState) -> GoalState:
|
||
result = self.run('goal.tactic', {"stateId": state.state_id, "goalId": 0, "conv": False})
|
||
if "error" in result:
|
||
raise ServerError(result["desc"])
|
||
if "tacticErrors" in result:
|
||
raise ServerError(result["tacticErrors"])
|
||
if "parseError" in result:
|
||
raise ServerError(result["parseError"])
|
||
return GoalState.parse(result)
|
||
|
||
|
||
def get_version():
|
||
import subprocess
|
||
with subprocess.Popen([_get_proc_path(), "--version"],
|
||
stdout=subprocess.PIPE,
|
||
cwd=_get_proc_cwd()) as p:
|
||
return p.communicate()[0].decode('utf-8').strip()
|
||
|
||
def extract_code_from_llm_output(reply):
|
||
i = reply.find("```lean")
|
||
if i != -1:
|
||
reply = reply[i + 7:]
|
||
i = reply.find("```")
|
||
reply = reply[:i]
|
||
return reply
|
||
i = reply.find("```")
|
||
if i != -1:
|
||
reply = reply[i + 3:]
|
||
i = reply.find("```")
|
||
reply = reply[:i]
|
||
return reply
|
||
return reply
|
||
|
||
class TestServer(unittest.TestCase):
|
||
|
||
def test_version(self):
|
||
self.assertEqual(get_version(), "0.2.14")
|
||
|
||
def test_goal_start(self):
|
||
server = Server()
|
||
state0 = server.goal_start("forall (p q: Prop), Or p q -> Or q p")
|
||
self.assertEqual(state0.state_id, 0)
|
||
state1 = server.goal_tactic(state0, goal_id=0, tactic="intro a")
|
||
self.assertEqual(state1.state_id, 1)
|
||
self.assertEqual(state1.goals, [Goal(
|
||
variables=[Variable(name="a", t="Prop")],
|
||
target="∀ (q : Prop), a ∨ q → q ∨ a",
|
||
name=None,
|
||
)])
|
||
self.assertEqual(str(state1.goals[0]),"a : Prop\n⊢ ∀ (q : Prop), a ∨ q → q ∨ a")
|
||
|
||
def test_conv_calc_sgl(self):
|
||
sgl.set_default_backend(sgl.OpenAI("gpt-4"))
|
||
|
||
server = Server()
|
||
state0 = server.goal_start("∀ (a b: Nat), (b = 2) -> 1 + a + 1 = a + b")
|
||
print("==========state0============")
|
||
print(state0)
|
||
variables = [
|
||
Variable(name="a", t="Nat"),
|
||
Variable(name="b", t="Nat"),
|
||
Variable(name="h", t="b = 2"),
|
||
]
|
||
|
||
state1 = server.goal_tactic(state0, goal_id=0, tactic="intro a b h")
|
||
print("==========state1============")
|
||
print(state1)
|
||
state2 = server.goal_tactic(state1, goal_id=0, tactic=TacticCalc("1 + a + 1 = a + 1 + 1"))
|
||
print("==========state2============")
|
||
print(state2)
|
||
self.assertEqual(state2.goals, [
|
||
Goal(
|
||
variables,
|
||
target="1 + a + 1 = a + 1 + 1",
|
||
name='calc',
|
||
),
|
||
Goal(
|
||
variables,
|
||
target="a + 1 + 1 = a + b",
|
||
),
|
||
])
|
||
state = select_tactic.run(str(state2))
|
||
for m in state.messages():
|
||
print(m["role"], ":", m["content"])
|
||
|
||
print("\n-- tactic --\n", state.stream_executor.variables)
|
||
print(state.stream_executor.arguments)
|
||
|
||
|
||
# print("==========state2============")
|
||
# print(state2)
|
||
# state_c1 = server.goal_conv_begin(state2, goal_id=0)
|
||
# print("==========state c1============")
|
||
# print(state_c1)
|
||
# state_c2 = server.goal_tactic(state_c1, goal_id=0, tactic="rhs")
|
||
# print("==========state c2============")
|
||
# print(state_c2)
|
||
# state_c3 = server.goal_tactic(state_c2, goal_id=0, tactic="rw [Nat.add_comm]")
|
||
# print("==========state c3============")
|
||
# print(state_c3)
|
||
# state_c4 = server.goal_conv_end(state_c3)
|
||
# print("==========state c4============")
|
||
# print(state_c4)
|
||
|
||
# state_c5 = server.goal_tactic(state_c4, goal_id=0, tactic="rfl")
|
||
# print("==========state c5============")
|
||
# print(state_c5)
|
||
# self.assertTrue(state_c5.is_solved)
|
||
|
||
# print()
|
||
|
||
# state3 = server.goal_tactic(state2, goal_id=1, tactic=TacticCalc("_ = a + 2"))
|
||
# print("==========state3============")
|
||
# print(state3)
|
||
# state4 = server.goal_tactic(state3, goal_id=0, tactic="rw [Nat.add_assoc]")
|
||
# print("==========state4============")
|
||
# print(state4)
|
||
# self.assertTrue(state4.is_solved)
|
||
|
||
|
||
def test_conv_calc(self):
|
||
server = Server()
|
||
state0 = server.goal_start("∀ (a b: Nat), (b = 2) -> 1 + a + 1 = a + b")
|
||
|
||
variables = [
|
||
Variable(name="a", t="Nat"),
|
||
Variable(name="b", t="Nat"),
|
||
Variable(name="h", t="b = 2"),
|
||
]
|
||
state1 = server.goal_tactic(state0, goal_id=0, tactic="intro a b h")
|
||
state2 = server.goal_tactic(state1, goal_id=0, tactic=TacticCalc("1 + a + 1 = a + 1 + 1"))
|
||
self.assertEqual(state2.goals, [
|
||
Goal(
|
||
variables,
|
||
target="1 + a + 1 = a + 1 + 1",
|
||
name='calc',
|
||
),
|
||
Goal(
|
||
variables,
|
||
target="a + 1 + 1 = a + b",
|
||
),
|
||
])
|
||
print("==========state2============")
|
||
print(state2)
|
||
state_c1 = server.goal_conv_begin(state2, goal_id=0)
|
||
print("==========state c1============")
|
||
print(state_c1)
|
||
state_c2 = server.goal_tactic(state_c1, goal_id=0, tactic="rhs")
|
||
print("==========state c2============")
|
||
print(state_c2)
|
||
state_c3 = server.goal_tactic(state_c2, goal_id=0, tactic="rw [Nat.add_comm]")
|
||
print("==========state c3============")
|
||
print(state_c3)
|
||
state_c4 = server.goal_conv_end(state_c3)
|
||
print("==========state c4============")
|
||
print(state_c4)
|
||
|
||
state_c5 = server.goal_tactic(state_c4, goal_id=0, tactic="rfl")
|
||
print("==========state c5============")
|
||
print(state_c5)
|
||
self.assertTrue(state_c5.is_solved)
|
||
|
||
print()
|
||
|
||
state3 = server.goal_tactic(state2, goal_id=1, tactic=TacticCalc("_ = a + 2"))
|
||
print("==========state3============")
|
||
print(state3)
|
||
state4 = server.goal_tactic(state3, goal_id=0, tactic="rw [Nat.add_assoc]")
|
||
print("==========state4============")
|
||
print(state4)
|
||
self.assertTrue(state4.is_solved)
|
||
|
||
def test_sglang_openai(self):
|
||
sgl.set_default_backend(sgl.OpenAI("gpt-4"))
|
||
|
||
print('\n----- Test sglang ---')
|
||
state = multi_turn_question.run(
|
||
question_1="What is the capital of the United States?",
|
||
question_2="List two local attractions.",
|
||
)
|
||
|
||
for m in state.messages():
|
||
print(m["role"], ":", m["content"])
|
||
|
||
print("\n-- answer_1 --\n", state["answer_1"])
|
||
|
||
|
||
if __name__ == '__main__':
|
||
|
||
unittest.main()
|
||
|