Merge pull request #35 from lean-dojo/dev

Refactor

.dockerignore

@@ -1,7 +1,6 @@
 *.txt
 .vscode
 build
-onnx*
 tmp
 lake-packages
 .lake

.gitignore

@@ -2,4 +2,4 @@
 /.lake
 
 .vscode
-*.olean
+*.olean

Dockerfile

@@ -1,6 +1,6 @@
 FROM ubuntu:latest
 
-WORKDIR /LeanInfer
+WORKDIR /LeanCopilot
 COPY . .
 
 # Install dependencies.
@@ -14,5 +14,5 @@ RUN curl https://raw.githubusercontent.com/leanprover/elan/master/elan-init.sh -
 
 # Build the Lean project.
 RUN lake build
-RUN lake script run LeanInfer/download
-RUN lake build LeanInferTests
+RUN lake script run LeanCopilot/download
+RUN lake build LeanCopilotTests

LeanCopilot.lean

@@ -0,0 +1,5 @@
+import LeanCopilot.Models
+import LeanCopilot.Frontend
+import LeanCopilot.Options
+import LeanCopilot.Tactics
+import LeanCopilot.LlmAesop

LeanCopilot/LlmAesop.lean

@@ -0,0 +1,23 @@
+import LeanCopilot.Tactics
+import LeanCopilot.Options
+import Aesop
+
+open Lean Lean.Elab.Command
+
+namespace LeanCopilot
+
+
+def tacGen : Aesop.TacGen := fun (mvarId : MVarId) => do
+ let state ← ppTacticState [mvarId]
+ let nm ← SuggestTactics.getGeneratorName
+ let model ← getGenerator nm
+ generate model state ""
+
+
+macro "#init_llm_aesop" : command => `(@[aesop 100%] def tacGen := LeanCopilot.tacGen)
+
+
+macro "search_proof" : tactic => `(tactic| aesop?)
+
+
+end LeanCopilot

LeanCopilot/Models.lean

@@ -0,0 +1,4 @@
+import LeanCopilot.Models.Interface
+import LeanCopilot.Models.Defs
+import LeanCopilot.Models.Registry
+import LeanCopilot.Models.FFI

LeanCopilot/Models/Builtin.lean

@@ -0,0 +1,27 @@
+import ModelCheckpointManager
+import LeanCopilot.Models.ByT5
+
+set_option autoImplicit false
+
+namespace LeanCopilot.Builtin
+
+
+def generator : NativeGenerator := {
+ url := Url.parse! "https://huggingface.co/kaiyuy/ct2-leandojo-lean4-tacgen-byt5-small"
+ tokenizer := ByT5.tokenizer
+ params := {
+ numReturnSequences := 32
+ }
+}
+
+
+def encoder : NativeEncoder := {
+ url := Url.parse! "https://huggingface.co/kaiyuy/ct2-leandojo-lean4-retriever-byt5-small"
+ tokenizer := ByT5.tokenizer
+}
+
+
+def premisesUrl := Url.parse! "https://huggingface.co/kaiyuy/premise-embeddings-leandojo-lean4-retriever-byt5-small"
+
+
+end LeanCopilot.Builtin

LeanInfer/Tokenization.lean → LeanCopilot/Models/ByT5.lean

@@ -1,8 +1,9 @@
-namespace LeanInfer
+import LeanCopilot.Models.Defs
 
-#eval "\t".length
+namespace LeanCopilot.ByT5
 
-def byt5Tokens : Array String := #[
+
+def vocab : Array String := #[
  "\u0000",
  "\u0001",
  "\u0002",
@@ -263,23 +264,29 @@ def byt5Tokens : Array String := #[
 
 
 def byteToToken (b : UInt8) : String :=
- byt5Tokens.get! b.toNat
+ vocab.get! b.toNat
 
 
 def tokenToByte! (t : String) : UInt8 :=
- byt5Tokens.findIdx? (· = t) |>.get! |>.toUInt8
+ vocab.findIdx? (· = t) |>.get! |>.toUInt8
 
 
-def tokenizeByt5 (text : String) (addEOS : Bool) : List String :=
- let tokens := byteToToken <$> text.toUTF8.toList
- if addEOS then
- tokens ++ ["</s>"]
- else
- tokens
+def tokenize (text : String) : Array String :=
+ (byteToToken <$> text.toUTF8.toList).toArray
 
 
-def detokenizeByt5 (tokens : Array String) : String :=
+def detokenize (tokens : Array String) : String :=
  String.fromUTF8Unchecked ⟨tokens.map tokenToByte!⟩
 
 
-end LeanInfer
+def eosToken := "</s>"
+
+
+def tokenizer : Tokenizer := {
+ tokenize := tokenize,
+ detokenize := detokenize,
+ eosToken := eosToken
+}
+
+
+end LeanCopilot.ByT5

LeanCopilot/Models/Defs.lean

@@ -0,0 +1,170 @@
+import Lean
+import ModelCheckpointManager
+import LeanCopilot.Models.Interface
+
+set_option autoImplicit false
+
+open Lean
+open System (FilePath)
+
+namespace LeanCopilot
+
+
+inductive Device where
+ | cpu
+ | cuda
+ | auto
+deriving Repr
+
+
+instance : Inhabited Device where
+ default := .auto
+
+
+def Device.toString : Device → String
+ | Device.cpu => "cpu"
+ | Device.cuda => "cuda"
+ | Device.auto => "auto"
+
+instance : ToString Device := ⟨Device.toString⟩
+
+
+inductive ComputeType where
+ | default
+ | auto
+ | int8
+ | int8_float32
+ | int8_float16
+ | int8_bfloat16
+ | int16
+ | float16
+ | bfloat16
+ | float32
+deriving Repr
+
+
+def ComputeType.toString : ComputeType → String
+ | ComputeType.default => "default"
+ | ComputeType.auto => "auto"
+ | ComputeType.int8 => "int8"
+ | ComputeType.int8_float32 => "int8_float32"
+ | ComputeType.int8_float16 => "int8_float16"
+ | ComputeType.int8_bfloat16 => "int8_bfloat16"
+ | ComputeType.int16 => "int16"
+ | ComputeType.float16 => "float16"
+ | ComputeType.bfloat16 => "bfloat16"
+ | ComputeType.float32 => "float32"
+
+
+instance : ToString ComputeType := ⟨ComputeType.toString⟩
+
+
+structure Tokenizer where
+ tokenize : String → Array String
+ detokenize : Array String → String
+ eosToken : String
+
+
+structure NativeModel where
+ url : Url
+ device : Device := .auto
+ deviceIndex : Array UInt64 := #[0]
+ computeType : ComputeType := .default
+ tokenizer : Tokenizer
+
+
+def NativeModel.name (model : NativeModel) : String := model.url.name!
+
+
+def NativeModel.path (model : NativeModel) : IO FilePath :=
+ getModelDir model.url
+
+
+structure BeamSearchParams where
+ numReturnSequences : UInt64
+ beamSize : UInt64 := numReturnSequences
+ minLength : UInt64 := 1
+ maxLength : UInt64 := 1024
+ lengthPenalty : Float := 0.0
+ patience : Float := 2.0
+ temperature : Float := 1.0
+deriving Repr
+
+
+structure NativeGenerator extends NativeModel where
+ params : BeamSearchParams
+
+
+structure NativeEncoder extends NativeModel
+
+
+structure ExternalModel where
+ name : String
+ host : String := "localhost"
+ port : UInt16 := 23333
+deriving Inhabited, Repr
+
+
+structure ExternalGenerator extends ExternalModel
+deriving Repr
+
+
+structure ExternalRequest where
+ name : String
+ input : String
+ «prefix» : String
+deriving ToJson
+
+
+structure ExternalResponse where
+ outputs : Array (String × Float)
+deriving FromJson
+
+
+def ExternalGenerator.generate (model : ExternalGenerator) (input : String) (targetPrefix : String) : IO $ Array (String × Float) := do
+ let url := s!"http://{model.host}:{model.port}/generate"
+ let req : ExternalRequest := {
+ name := model.name,
+ input := input,
+ «prefix» := targetPrefix
+ }
+ let reqStr := (toJson req).pretty 99999999999999999
+ let out ← IO.Process.run {
+ cmd := "curl"
+ args := #["-X", "POST", url, "-H", "accept: application/json", "-H", "Content-Type: application/json", "-d", reqStr]
+ }
+
+ let some json := Json.parse out |>.toOption | throw $ IO.userError "Failed to parse response"
+ let some res := (fromJson? json : Except String ExternalResponse) |>.toOption | throw $ IO.userError "Failed to parse response"
+ return res.outputs
+
+
+instance : TextToText ExternalGenerator := ⟨ExternalGenerator.generate⟩
+
+
+structure ExternalEncoder extends ExternalModel
+deriving Repr
+
+
+def ExternalEncoder.encode (model : ExternalEncoder) (input : String) : IO FloatArray := do
+ return FloatArray.mk #[0.0]
+
+
+instance : TextToVec ExternalEncoder := ⟨ExternalEncoder.encode⟩
+
+
+structure GenericGenerator where
+ generate : String → String → IO (Array (String × Float))
+
+
+instance : TextToText GenericGenerator := ⟨GenericGenerator.generate⟩
+
+
+structure GenericEncoder where
+ encode : String → IO FloatArray
+
+
+instance : TextToVec GenericEncoder := ⟨GenericEncoder.encode⟩
+
+
+end LeanCopilot