237 lines
8.9 KiB
Python
237 lines
8.9 KiB
Python
from pantograph.server import Server, ServerError, TacticFailure
|
||
from pantograph.expr import Variable, Goal, TacticCalc
|
||
import unittest
|
||
import sglang as sgl
|
||
|
||
LEAN4_INTRO = '''/-- A sequence `u` of real numbers converges to `l` if `∀ ε > 0, ∃ N, ∀ n ≥ N, |u_n - l| ≤ ε`.
|
||
This condition will be spelled `seq_limit u l`. -/
|
||
def seq_limit (u : ℕ → ℝ) (l : ℝ) : Prop :=
|
||
∀ ε > 0, ∃ N, ∀ n ≥ N, |u n - l| ≤ ε
|
||
|
||
/- In the above definition, note that the `n`-th term of the sequence `u` is denoted
|
||
simply by `u n`.
|
||
|
||
Similarly, in the next definition, `f x` is what we would write `f(x)` on paper.
|
||
Also note that implication is denoted by a single arrow (we'll explain why later). -/
|
||
|
||
/-- A function`f : ℝ → ℝ` is continuous at `x₀` if
|
||
`∀ ε > 0, ∃ δ > 0, ∀ x, |x - x₀| ≤ δ ⇒ |f(x) - f(x₀)| ≤ ε`.
|
||
This condition will be spelled `continuous_at f x₀`.-/
|
||
def continuous_at (f : ℝ → ℝ) (x₀ : ℝ) : Prop :=
|
||
∀ ε > 0, ∃ δ > 0, ∀ x, |x - x₀| ≤ δ → |f x - f x₀| ≤ ε
|
||
|
||
/-- Now we claim that if `f` is continuous at `x₀` then it is sequentially continuous
|
||
at `x₀`: for any sequence `u` converging to `x₀`, the sequence `f ∘ u` converges
|
||
to `f x₀`. -/
|
||
example (f : ℝ → ℝ) (u : ℕ → ℝ) (x₀ : ℝ) (hu : seq_limit u x₀) (hf : continuous_at f x₀) :
|
||
seq_limit (f ∘ u) (f x₀) := by { -- This `by` keyword marks the beginning of the proof
|
||
-- Put your text cursor here and watch the Lean InfoView panel to the right.
|
||
-- Then move your cursor from line to line in the proof while monitoring the Infoview.
|
||
|
||
-- Our goal is to prove that, for any positive `ε`, there exists a natural
|
||
-- number `N` such that, for any natural number `n` at least `N`,
|
||
-- `|f(u_n) - f(x₀)|` is at most `ε`.
|
||
unfold seq_limit
|
||
-- Fix a positive number `ε`.
|
||
intros ε hε
|
||
-- By assumption on `f` applied to this positive `ε`, we get a positive `δ`
|
||
-- such that, for all real number `x`, if `|x - x₀| ≤ δ` then `|f(x) - f(x₀)| ≤ ε` (1).
|
||
obtain ⟨δ, δ_pos, Hf⟩ : ∃ δ > 0, ∀ x, |x - x₀| ≤ δ → |f x - f x₀| ≤ ε := hf ε hε
|
||
-- The assumption on `u` applied to this `δ` gives a natural number `N` such that
|
||
-- for every natural number `n`, if `n ≥ N` then `|u_n - x₀| ≤ δ` (2).
|
||
obtain ⟨N, Hu⟩ : ∃ N, ∀ n ≥ N, |u n - x₀| ≤ δ := hu δ δ_pos
|
||
-- Let's prove `N` is suitable.
|
||
use N
|
||
-- Fix `n` which is at least `N`. Let's prove `|f(u_n) - f(x₀)| ≤ ε`.
|
||
intros n hn
|
||
-- Thanks to (1) applied to `u_n`, it suffices to prove that `|u_n - x₀| ≤ δ`.
|
||
apply Hf
|
||
-- This follows from property (2) and our assumption on `n`.
|
||
exact Hu n hn
|
||
-- This finishes the proof!
|
||
}
|
||
|
||
/-
|
||
Now that this proof is over, you can use the file explorer to the
|
||
left of this panel to open the file `Exercises > 01Rewriting.lean`.
|
||
-/'''
|
||
|
||
LEAN4_REWRITE = '''Rewrite tactic tutorial:
|
||
example (a b c : Nat) : a + b + c = a + c + b := by
|
||
rw [Nat.add_assoc, Nat.add_comm b, ← Nat.add_assoc]
|
||
|
||
example (a b c : Nat) : a + b + c = a + c + b := by
|
||
rw [Nat.add_assoc, Nat.add_assoc, Nat.add_comm b]
|
||
|
||
example (a b c : Nat) : a + b + c = a + c + b := by
|
||
rw [Nat.add_assoc, Nat.add_assoc, Nat.add_comm _ b]
|
||
|
||
example (f : Nat → Nat) (a : Nat) (h : a + 0 = 0) : f a = f 0 := by
|
||
rw [Nat.add_zero] at h
|
||
rw [h]
|
||
|
||
def Tuple (α : Type) (n : Nat) :=
|
||
{ as : List α // as.length = n }
|
||
|
||
example (n : Nat) (h : n = 0) (t : Tuple α n) : Tuple α 0 := by
|
||
rw [h] at t
|
||
exact t
|
||
'''
|
||
|
||
@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, server, state, goal_id, feedback_turns = 5):
|
||
|
||
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(LEAN4_REWRITE)
|
||
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))
|
||
for i in range(feedback_turns):
|
||
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+"```")
|
||
success, new_state = apply_tactic(server, state, goal_id, tactic)
|
||
# print("===execute===")
|
||
# print(success, new_state )
|
||
if not success:
|
||
with s.user():
|
||
s += "This answer got Lean compile error:\n" + str(new_state) + "\n"
|
||
s += "Please try again by taking the Lean compiler feedback."
|
||
|
||
else:
|
||
return tactic, new_state
|
||
return None, None
|
||
|
||
def apply_tactic(server, state, goal_id, tactic):
|
||
try:
|
||
new_state = server.goal_tactic(state=state, goal_id=goal_id, tactic=tactic)
|
||
except ServerError as e:
|
||
return False, e
|
||
except TacticFailure as e:
|
||
return False, e
|
||
return True, new_state
|
||
|
||
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 TestServerSGL(unittest.TestCase):
|
||
|
||
def test_conv_calc_sgl(self):
|
||
n_trails = 5
|
||
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",
|
||
),
|
||
])
|
||
state3 = None
|
||
for i in range(n_trails):
|
||
print(f"===============trail {str(i)}============")
|
||
try:
|
||
state = select_tactic.run(server, state2, goal_id = 1)
|
||
state3 = state.ret_value
|
||
for m in state.messages():
|
||
print(m["role"], ":", m["content"])
|
||
|
||
print("\n-- new state --\n", state3)
|
||
break
|
||
|
||
except ServerError as e:
|
||
print(f"server error: {e}")
|
||
continue
|
||
state3 = server.goal_tactic(state2, goal_id=1, tactic=TacticCalc("_ = a + 2"))
|
||
|
||
|
||
print("==========state3============")
|
||
print(state3)
|
||
state4 = None
|
||
for i in range(n_trails):
|
||
print(f"===============trail {str(i)}============")
|
||
try:
|
||
state = select_tactic.run(server, state3, goal_id = 0)
|
||
state4 = state.ret_value
|
||
for m in state.messages():
|
||
print(m["role"], ":", m["content"])
|
||
|
||
print("\n-- new state --\n", state4)
|
||
break
|
||
|
||
except ServerError as e:
|
||
print(f"server error: {e}")
|
||
continue
|
||
|
||
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()
|
||
|