From 7f5b9ae6f3c30d6283794a7bc176a8fc534447e2 Mon Sep 17 00:00:00 2001
From: Kaiyu Yang <kaiyuy@caltech.edu>
Date: Mon, 27 Nov 2023 19:25:16 -0800
Subject: [PATCH 1/4] use openblas

---
 LeanInfer/Cache.lean         |   4 +-
 LeanInferTests/Examples.lean |   5 +
 README.md                    |   3 +
 lakefile.lean                | 282 ++++++++++++++++++++---------------
 scripts/convert_to_ct2.py    |   2 +-
 scripts/convert_to_onnx.py   |  12 +-
 6 files changed, 184 insertions(+), 124 deletions(-)

diff --git a/LeanInfer/Cache.lean b/LeanInfer/Cache.lean
index 0f0a8d0..67fef74 100644
--- a/LeanInfer/Cache.lean
+++ b/LeanInfer/Cache.lean
@@ -11,7 +11,7 @@ private def getHomeDir : IO FilePath := do
   return dir
 
 
-private def ensureExists (dir : FilePath) : IO Unit := do
+private def ensureDirExists (dir : FilePath) : IO Unit := do
   if !(← dir.pathExists)  then
     IO.FS.createDirAll dir
 
@@ -25,7 +25,7 @@ def getCacheDir : IO FilePath := do
   let dir := match ← IO.getEnv "LEAN_INFER_CACHE_DIR" with
   | some dir => (dir : FilePath)
   | none => defaultCacheDir
-  ensureExists dir
+  ensureDirExists dir
   return dir.normalize
 
 
diff --git a/LeanInferTests/Examples.lean b/LeanInferTests/Examples.lean
index 9f73c72..ff8d33a 100644
--- a/LeanInferTests/Examples.lean
+++ b/LeanInferTests/Examples.lean
@@ -33,3 +33,8 @@ example (a b c : Nat) : a + b + c = a + c + b := by
 example (a b c : Nat) : a + b + c = a + c + b := by
   suggest_tactics "rw"  -- You may provide a prefix to constrain the generated tactics.
   sorry
+
+
+example (a b c : Nat) : a + b + c = a + c + b := by
+  select_premises
+  sorry
diff --git a/README.md b/README.md
index 37bd406..48baba2 100644
--- a/README.md
+++ b/README.md
@@ -9,6 +9,9 @@ LeanInfer provides tactic suggestions by running LLMs through Lean's foreign fun
 ## Requirements
 
 * Supported platforms: Linux and macOS (:warning: maybe also Windows WSL, but untested)
+* xcode-select --install
+
+
 * Git LFS
 * A C++17 compatible compiler, e.g., recent versions of GCC or Clang
 * CMake >= 3.7
diff --git a/lakefile.lean b/lakefile.lean
index 3a78385..8c76a3a 100644
--- a/lakefile.lean
+++ b/lakefile.lean
@@ -9,7 +9,7 @@ inductive SupportedOS where
 deriving Inhabited, BEq
 
 
-def getOS : IO SupportedOS := do
+def getOS! : IO SupportedOS := do
   if Platform.isWindows then
     error "Windows is not supported"
   if Platform.isOSX then
@@ -24,35 +24,55 @@ inductive SupportedArch where
 deriving Inhabited, BEq
 
 
-def getArch? : BaseIO (Option SupportedArch) := do
-  return some .x86_64
+def nproc : IO Nat := do
+  let out ← IO.Process.output {cmd := "nproc", stdin := .null}
+  return out.stdout.trim.toNat!
 
 
-def getArch : IO SupportedArch := do
+def getArch? : IO (Option SupportedArch) := do
+  let out ← IO.Process.output {cmd := "uname", args := #["-m"], stdin := .null}
+  let arch := out.stdout.trim
+  if arch ∈ ["arm64", "aarch64"] then
+    return some .arm64
+  else if arch == "x86_64" then
+    return some .x86_64
+  else
+    return none
+
+
+def getArch! : IO SupportedArch := do
   if let some arch ← getArch? then
     return arch
   else
     error "Unknown architecture"
 
 
+def isArm! : IO Bool := do
+  return (← getArch!) == .arm64
+
+
+def isArm : Bool := run_io isArm!
+
+
 structure SupportedPlatform where
   os : SupportedOS
   arch : SupportedArch
 
 
-def getPlatform : IO SupportedPlatform := do
+def getPlatform! : IO SupportedPlatform := do
   if Platform.numBits != 64 then
     error "Only 64-bit platforms are supported"
-  return ⟨← getOS, ← getArch⟩
+  return ⟨← getOS!, ← getArch!⟩
 
 
 package LeanInfer where
-  -- preferReleaseBuild := get_config? noCloudRelease |>.isNone
-  -- buildArchive? := is_arm? |>.map (if · then "arm64" else "x86_64")
+  preferReleaseBuild := get_config? noCloudRelease |>.isNone
+  buildArchive? := if isArm then "arm64" else "x86_64"
   precompileModules := true
-  buildType := BuildType.release
+  buildType := BuildType.debug  -- TODO: Change to release.
   moreLinkArgs := #[s!"-L{__dir__}/.lake/build/lib", "-lonnxruntime", "-lctranslate2"]
-  weakLeanArgs := #[s!"--load-dynlib={__dir__}/.lake/build/lib/" ++ nameToSharedLib "onnxruntime", s!"--load-dynlib={__dir__}/.lake/build/lib/" ++ nameToSharedLib "ctranslate2"]
+  weakLeanArgs := #["onnxruntime", "ctranslate2"].map fun name =>
+    s!"--load-dynlib={__dir__}/.lake/build/lib/" ++ nameToSharedLib name
 
 
 @[default_target]
@@ -85,22 +105,32 @@ def afterReleaseAsync (pkg : Package) (build : BuildM α) : IndexBuildM (Job α)
     Job.async build
 
 
-def getOnnxPlatform : IO String := do
-  let ⟨os, arch⟩  ← getPlatform
+def getOnnxPlatform! : IO String := do
+  let ⟨os, arch⟩  ← getPlatform!
   match os with
   | .linux => return if arch == .x86_64 then "linux-x64" else "linux-aarch64"
   | .macos => return "osx-universal2"
 
 
+def hasCUDA : IO Bool := do
+  let out ← IO.Process.output {cmd := "which", args := #["nvcc"], stdin := .null}
+  return out.exitCode == 0
+
+
+def ensureDirExists (dir : FilePath) : IO Unit := do
+  if !(← dir.pathExists)  then
+    IO.FS.createDirAll dir
+
+
 /- Download and Copy ONNX's C++ header files to `build/include` and shared libraries to `build/lib` -/
 target libonnxruntime pkg : FilePath := do
   afterReleaseAsync pkg do
-  let _ ← getPlatform
+  let _ ← getPlatform!
   let dst := pkg.nativeLibDir / (nameToSharedLib "onnxruntime")
   createParentDirs dst
 
   let onnxVersion := "1.15.1"
-  let onnxFileStem := s!"onnxruntime-{← getOnnxPlatform}-{onnxVersion}"
+  let onnxFileStem := s!"onnxruntime-{← getOnnxPlatform!}-{onnxVersion}"
   let onnxFilename := onnxFileStem ++ ".tgz"
   let onnxURL := "https://github.com/microsoft/onnxruntime/releases/download/v1.15.1/" ++ onnxFilename
 
@@ -122,9 +152,18 @@ target libonnxruntime pkg : FilePath := do
         cmd := "cp"
         args := #[dst.toString, dst'.toString]
       }
+      ensureDirExists $ pkg.buildDir / "include"
       proc {
         cmd := "cp"
-        args := #["-r", (onnxStem / "include").toString, (pkg.buildDir / "include").toString]
+        args := #[(onnxStem / "include" / "onnxruntime_cxx_api.h").toString, (pkg.buildDir / "include").toString ++ "/"]
+      }
+      proc {
+        cmd := "cp"
+        args := #[(onnxStem / "include" / "onnxruntime_cxx_inline.h").toString, (pkg.buildDir / "include").toString ++ "/"]
+      }
+      proc {
+        cmd := "cp"
+        args := #[(onnxStem / "include" / "onnxruntime_c_api.h").toString, (pkg.buildDir / "include").toString ++ "/"]
       }
       proc {
         cmd := "rm"
@@ -136,140 +175,150 @@ target libonnxruntime pkg : FilePath := do
 
 
 def gitClone (url : String) (cwd : Option FilePath) : LogIO Unit := do
-  proc {
+  proc (quiet := true) {
     cmd := "git"
     args := #["clone", "--recursive", url]
     cwd := cwd
   }
 
 
-def runCmake (root : FilePath) (flags : Array String) : LogIO Bool := do
+def runCmake (root : FilePath) (flags : Array String) : LogIO Unit := do
   assert! (← root.pathExists) ∧ (← (root / "CMakeLists.txt").pathExists)
   let buildDir := root / "build"
   if ← buildDir.pathExists then
     IO.FS.removeDirAll buildDir
   IO.FS.createDirAll buildDir
-  testProc {
+  let ok ← testProc {
     cmd := "cmake"
     args := flags ++ #[".."]
     cwd := buildDir
   }
+  if ¬ ok then
+    error "Failed to run cmake"
 
 
-def autoCt2Cmake (root : FilePath) : LogIO Unit := do
-  let basicFlags := #["-DBUILD_CLI=OFF", "-DOPENMP_RUNTIME=NONE", "-DWITH_CUDA=OFF", "-DWITH_CUDNN=OFF", "-DWITH_DNNL=OFF", "-DWITH_MKL=OFF", "-DWITH_ACCELERATE=OFF", "-DWITH_OPENBLAS=OFF"]
-  assert! ← runCmake root basicFlags
-
-  let hasOpenMP ← runCmake root (basicFlags.erase "-DOPENMP_RUNTIME=NONE" |>.push "-DOPENMP_RUNTIME=COMP")
-  let hasCUDA ← runCmake root (basicFlags.erase "-DWITH_CUDA=OFF" |>.push "-DWITH_CUDA=ON")
-  let hasCuDNN := hasCUDA && (← runCmake root ((basicFlags.erase "-DWITH_CUDA=OFF" |>.erase "-DWITH_CUDNN=OFF") ++ #["-DWITH_CUDA=ON", "-DWITH_CUDNN=ON"]))
-  let hasDNNL ← runCmake root (basicFlags.erase "-DWITH_DNNL=OFF" |>.push "-DWITH_DNNL=ON")
-  let hasAccelerate := Platform.isOSX && (← runCmake root (basicFlags.erase "-DWITH_ACCELERATE=OFF" |>.push "-DWITH_ACCELERATE=ON"))
-  let hasMKL := ¬ hasAccelerate && (← runCmake root (basicFlags.erase "-DWITH_MKL=OFF" |>.push "-DWITH_MKL=ON"))
-  let hasOpenBLAS := ¬ hasDNNL && ¬ hasAccelerate && ¬ hasMKL && (← runCmake root (basicFlags.erase "-DWITH_OPENBLAS=OFF" |>.push "-DWITH_OPENBLAS=ON"))
-
-  if ¬ (hasMKL ∨ hasAccelerate ∨ hasDNNL ∨ hasOpenBLAS) then
-    error "No BLAS library found. You need at least one of MKL, Accelerate, DNNL, or OpenBLAS. See https://github.com/lean-dojo/LeanInfer#requirements for details."
-
-  let flags := #[
-    "-DBUILD_CLI=OFF",
-    "-DOPENMP_RUNTIME=" ++ (if hasOpenMP then "COMP" else "NONE"),
-    "-DWITH_CUDA=" ++ (if hasCUDA then "ON" else "OFF"),
-    "-DWITH_CUDNN=" ++ (if hasCuDNN then "ON" else "OFF"),
-    "-DWITH_DNNL=" ++ (if hasDNNL then "ON" else "OFF"),
-    "-DWITH_MKL=" ++ (if hasMKL then "ON" else "OFF"),
-    "-DWITH_ACCELERATE=" ++ (if hasAccelerate then "ON" else "OFF"),
-    "-DWITH_OPENBLAS=" ++ (if hasOpenBLAS then "ON" else "OFF")
-    ]
-  logInfo s!"Using CTranslate2 CMake flags: {flags}"
-  assert! ← runCmake root flags
-
-/--
-Build CTranslate2 from source using cmake.
-TODO: Include the flags into the trace.
--/
-def buildCmakeProject (root : FilePath) : LogIO Unit := do
-  -- Run cmake.
-  if let some flags := get_config? Ct2Flags then
-    logInfo s!"Using CTranslate2 cmake flags: {flags}"
-    let _ ← runCmake root flags.splitOn.toArray
+target libopenblas pkg : FilePath := do
+  afterReleaseAsync pkg do
+    let rootDir := pkg.buildDir / "OpenBLAS"
+    ensureDirExists rootDir
+    let dst := rootDir / (nameToSharedLib "openblas")
+    let url := "https://github.com/OpenMathLib/OpenBLAS"
+
+    try
+      let depTrace := Hash.ofString url
+      let trace ← buildFileUnlessUpToDate dst depTrace do
+        logStep s!"Cloning OpenBLAS from {url}"
+        gitClone url pkg.buildDir
+
+        let numThreads := min 32 (← nproc)
+        let flags := #["NO_LAPACK=1", "NO_FORTRAN=1", s!"-j{numThreads}"]
+        logStep s!"Building OpenBLAS with `make{flags.foldl (· ++ " " ++ ·) ""}`"
+        proc (quiet := true) {
+          cmd := "make"
+          args := flags
+          cwd := rootDir
+        }
+      return (dst, trace)
+
+    else
+      return (dst, ← computeTrace dst)
+
+
+def getCt2CmakeFlags : IO (Array String) := do
+  let mut flags := #["-DBUILD_CLI=OFF", "-DOPENMP_RUNTIME=NONE", "-DWITH_DNNL=OFF", "-DWITH_MKL=OFF"]
+
+  match ← getOS! with
+  | .macos => flags := flags ++ #["-DWITH_ACCELERATE=ON", "-DWITH_OPENBLAS=OFF"]
+  | .linux => flags := flags ++ #["-DWITH_ACCELERATE=OFF", "-DWITH_OPENBLAS=ON", "-DOPENBLAS_INCLUDE_DIR=../../OpenBLAS", "-DOPENBLAS_LIBRARY=../../OpenBLAS/libopenblas.so"]
+
+  if ← hasCUDA then
+    flags := flags ++ #["-DWITH_CUDA=ON", "-DWITH_CUDNN=ON"]
   else
-    autoCt2Cmake root
-  -- Run make.
-  proc {
-    cmd := "make"
-    args := #["-j8"]
-    cwd := root / "build"
-  }
+    flags := flags ++ #["-DWITH_CUDA=OFF", "-DWITH_CUDNN=OFF"]
+
+  return flags
 
 
 /- Download and build CTranslate2. Copy its C++ header files to `build/include` and shared libraries to `build/lib` -/
 target libctranslate2 pkg : FilePath := do
-  afterReleaseAsync pkg do
-  let _ ← getPlatform
-  let dst := pkg.nativeLibDir / (nameToSharedLib "ctranslate2")
-  createParentDirs dst
-  let ct2URL := "https://github.com/OpenNMT/CTranslate2"
-
-  try
-    let depTrace := Hash.ofString ct2URL
-    let trace ← buildFileUnlessUpToDate dst depTrace do
-      logStep s!"Cloning CTranslate2 from {ct2URL} into {pkg.buildDir}"
-      gitClone ct2URL pkg.buildDir
-
-      logStep s!"Building CTranslate2 with cmake"
-      let ct2Dir := pkg.buildDir / "CTranslate2"
-      buildCmakeProject ct2Dir
-      proc {
-        cmd := "cp"
-        args := #[(ct2Dir / "build" / nameToSharedLib "ctranslate2").toString, dst.toString]
-      }
-      -- TODO: Don't hardcode the version "3".
-      let dst' := pkg.nativeLibDir / (nameToVersionedSharedLib "ctranslate2" "3")
-      proc {
-        cmd := "cp"
-        args := #[dst.toString, dst'.toString]
-      }
-      proc {
-        cmd := "cp"
-        args := #["-r", (ct2Dir / "include" / "ctranslate2").toString, (pkg.buildDir / "include" / "ctranslate2").toString]
-      }
-      proc {
-        cmd := "cp"
-        args := #["-r", (ct2Dir / "include" / "nlohmann").toString, (pkg.buildDir / "include" / "nlohmann").toString]
-      }
-      proc {
-        cmd := "cp"
-        args := #["-r", (ct2Dir / "include" / "half_float").toString, (pkg.buildDir / "include" / "half_float").toString]
-      }
-      proc {
-        cmd := "rm"
-        args := #["-rf", ct2Dir.toString]
-      }
-    return (dst, trace)
-  else
-    return (dst, ← computeTrace dst)
-
+  if (← getOS!) == .linux then
+    let openblas ← libopenblas.fetch
+    let _ ← openblas.await
 
-def buildCpp (pkg : Package) (path : FilePath) (deps : List (BuildJob FilePath)) : SchedulerM (BuildJob FilePath) := do
+  afterReleaseAsync pkg do
+    let dst := pkg.nativeLibDir / (nameToSharedLib "ctranslate2")
+    createParentDirs dst
+    let ct2URL := "https://github.com/OpenNMT/CTranslate2"
+
+    try
+      let depTrace := Hash.ofString ct2URL
+      let trace ← buildFileUnlessUpToDate dst depTrace do
+        logStep s!"Cloning CTranslate2 from {ct2URL}"
+        gitClone ct2URL pkg.buildDir
+
+        let ct2Dir := pkg.buildDir / "CTranslate2"
+        let flags ← getCt2CmakeFlags
+        logStep s!"Configuring CTranslate2 with `cmake{flags.foldl (· ++ " " ++ ·) ""} ..`"
+        runCmake ct2Dir flags
+        let numThreads := min 32 (← nproc)
+        logStep s!"Building CTranslate2 with `make -j{numThreads}`"
+        proc {
+          cmd := "make"
+          args := #[s!"-j{numThreads}"]
+          cwd := ct2Dir / "build"
+        }
+
+        ensureDirExists $ pkg.buildDir / "include"
+        proc {
+          cmd := "cp"
+          args := #[(ct2Dir / "build" / nameToSharedLib "ctranslate2").toString, dst.toString]
+        }
+        -- TODO: Don't hardcode the version "3".
+        let dst' := pkg.nativeLibDir / (nameToVersionedSharedLib "ctranslate2" "3")
+        proc {
+          cmd := "cp"
+          args := #[dst.toString, dst'.toString]
+        }
+        proc {
+          cmd := "cp"
+          args := #["-r", (ct2Dir / "include" / "ctranslate2").toString, (pkg.buildDir / "include" / "ctranslate2").toString]
+        }
+        proc {
+          cmd := "cp"
+          args := #["-r", (ct2Dir / "include" / "nlohmann").toString, (pkg.buildDir / "include" / "nlohmann").toString]
+        }
+        proc {
+          cmd := "cp"
+          args := #["-r", (ct2Dir / "include" / "half_float").toString, (pkg.buildDir / "include" / "half_float").toString]
+        }
+        proc {
+          cmd := "rm"
+          args := #["-rf", ct2Dir.toString]
+        }
+      return (dst, trace)
+    else
+      return (dst, ← computeTrace dst)
+
+
+def buildCpp (pkg : Package) (path : FilePath) (dep : BuildJob FilePath) : SchedulerM (BuildJob FilePath) := do
   let optLevel := if pkg.buildType == .release then "-O3" else "-O0"
   let flags := #["-fPIC", "-std=c++17", optLevel]
   let args := flags ++ #["-I", (← getLeanIncludeDir).toString, "-I", (pkg.buildDir / "include").toString]
   let oFile := pkg.buildDir / (path.withExtension "o")
   let srcJob ← inputFile <| pkg.dir / path
-  buildFileAfterDepList oFile (srcJob :: deps) (extraDepTrace := computeHash flags) fun deps =>
+  buildFileAfterDepList oFile [srcJob, dep] (extraDepTrace := computeHash flags) fun deps =>
     compileO path.toString oFile deps[0]! args "c++"
 
 
 target onnx.o pkg : FilePath := do
   let onnx ← libonnxruntime.fetch
-  let build := buildCpp pkg "cpp/onnx.cpp" [onnx]
+  let build := buildCpp pkg "cpp/onnx.cpp" onnx
   afterReleaseSync pkg build
 
 
 target ct2.o pkg : FilePath := do
   let ct2 ← libctranslate2.fetch
-  let build := buildCpp pkg "cpp/ct2.cpp" [ct2]
+  let build := buildCpp pkg "cpp/ct2.cpp" ct2
   afterReleaseSync pkg build
 
 
@@ -300,21 +349,18 @@ def initGitLFS : IO Unit := do
 
 def HF_BASE_URL := "https://huggingface.co"
 
+
 structure HuggingFaceURL where
   user : Option String
   modelName : String
 
+
 instance : ToString HuggingFaceURL where
   toString url := match url.user with
   | none => s!"{HF_BASE_URL}/{url.modelName}"
   | some user => s!"{HF_BASE_URL}/{user}/{url.modelName}"
 
 
-def ensureExists (dir : FilePath) : IO Unit := do
-  if !(← dir.pathExists)  then
-    IO.FS.createDirAll dir
-
-
 def getHomeDir : IO FilePath := do
   let some dir ← IO.getEnv "HOME" | throw $ IO.userError "Cannot find the $HOME environment variable."
   return dir
@@ -329,7 +375,7 @@ def getCacheDir : IO FilePath := do
   let dir := match ← IO.getEnv "LEAN_INFER_CACHE_DIR" with
   | some dir => (dir : FilePath)
   | none => defaultCacheDir
-  ensureExists dir
+  ensureDirExists dir
   return dir.normalize
 
 
@@ -361,7 +407,7 @@ def downloadIfNecessary (url : HuggingFaceURL) : IO Unit := do
 
   println! s!"Downloading the model into {dir}"
   let some parentDir := dir.parent | unreachable!
-  ensureExists parentDir
+  ensureDirExists parentDir
   initGitLFS
   let proc ← IO.Process.output {
     cmd := "git"
diff --git a/scripts/convert_to_ct2.py b/scripts/convert_to_ct2.py
index fcfea4d..50b1c50 100644
--- a/scripts/convert_to_ct2.py
+++ b/scripts/convert_to_ct2.py
@@ -124,4 +124,4 @@ def set_layer_norm(self, spec, layer_norm):
 output = encoder.forward_batch(
     [tokenizer.convert_ids_to_tokens(tokenizer.encode(state))]
 )
-feature = np.array(output).mean(axis=1)
+feature = np.array(output.last_hidden_state).mean(axis=1)
diff --git a/scripts/convert_to_onnx.py b/scripts/convert_to_onnx.py
index f8be388..2c59216 100644
--- a/scripts/convert_to_onnx.py
+++ b/scripts/convert_to_onnx.py
@@ -1,5 +1,11 @@
+import os
 from optimum.onnxruntime import ORTModelForSeq2SeqLM
 
-model_name = "leandojo-lean4-tacgen-byt5-small"
-model = ORTModelForSeq2SeqLM.from_pretrained(f"kaiyuy/{model_name}", export=True)
-model.save_pretrained(f"onnx-{model_name}")
+model_name = "leandojo-lean4-retriever-byt5-small"
+ORTModelForSeq2SeqLM.from_pretrained(
+    f"kaiyuy/{model_name}", export=True
+).save_pretrained(f"onnx-{model_name}")
+
+os.remove(f"onnx-{model_name}/decoder_model.onnx")
+os.remove(f"onnx-{model_name}/decoder_with_past_model.onnx")
+os.remove(f"onnx-{model_name}/generation_config.json")
\ No newline at end of file

From ea1d2be24e3a800eb8b7e793c385a00129dd9e85 Mon Sep 17 00:00:00 2001
From: Kaiyu Yang <kaiyuy@caltech.edu>
Date: Tue, 28 Nov 2023 06:51:54 -0800
Subject: [PATCH 2/4] clean

---
 Dockerfile               |  2 +-
 lakefile.lean            | 21 +++++++++++++--------
 scripts/build.sh         |  3 ---
 scripts/build_example.sh |  3 ---
 4 files changed, 14 insertions(+), 15 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index 702d093..7d0b7b3 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -4,7 +4,7 @@ WORKDIR /LeanInfer
 COPY . .
 
 # Install dependencies.
-RUN apt-get update && apt-get install -y curl wget git git-lfs cmake clang lld libc++-dev libopenblas-dev
+RUN apt-get update && apt-get install -y curl wget git git-lfs cmake clang lld libc++-dev
 RUN git lfs update --force
 
 # Install elan.
diff --git a/lakefile.lean b/lakefile.lean
index 8c76a3a..f28560f 100644
--- a/lakefile.lean
+++ b/lakefile.lean
@@ -51,7 +51,17 @@ def isArm! : IO Bool := do
   return (← getArch!) == .arm64
 
 
-def isArm : Bool := run_io isArm!
+def hasCUDA : IO Bool := do
+  let out ← IO.Process.output {cmd := "which", args := #["nvcc"], stdin := .null}
+  return out.exitCode == 0
+
+
+def buildArchiveName : String :=
+  let arch := if run_io isArm! then "arm64" else "x86_64"
+  if run_io hasCUDA then
+    s!"{arch}-cuda"
+  else
+    arch
 
 
 structure SupportedPlatform where
@@ -67,9 +77,9 @@ def getPlatform! : IO SupportedPlatform := do
 
 package LeanInfer where
   preferReleaseBuild := get_config? noCloudRelease |>.isNone
-  buildArchive? := if isArm then "arm64" else "x86_64"
+  buildArchive? := buildArchiveName
   precompileModules := true
-  buildType := BuildType.debug  -- TODO: Change to release.
+  buildType := BuildType.release
   moreLinkArgs := #[s!"-L{__dir__}/.lake/build/lib", "-lonnxruntime", "-lctranslate2"]
   weakLeanArgs := #["onnxruntime", "ctranslate2"].map fun name =>
     s!"--load-dynlib={__dir__}/.lake/build/lib/" ++ nameToSharedLib name
@@ -112,11 +122,6 @@ def getOnnxPlatform! : IO String := do
   | .macos => return "osx-universal2"
 
 
-def hasCUDA : IO Bool := do
-  let out ← IO.Process.output {cmd := "which", args := #["nvcc"], stdin := .null}
-  return out.exitCode == 0
-
-
 def ensureDirExists (dir : FilePath) : IO Unit := do
   if !(← dir.pathExists)  then
     IO.FS.createDirAll dir
diff --git a/scripts/build.sh b/scripts/build.sh
index 50caacd..50e89ae 100644
--- a/scripts/build.sh
+++ b/scripts/build.sh
@@ -4,9 +4,6 @@
 # 1. Launch a codespace for LeanInfer.
 # 2. Run `source scripts/build.sh`.
 
-# Install OpenBLAS.
-sudo apt-get update && sudo apt-get install -y libopenblas-dev
-
 # Set up elan.
 curl https://raw.githubusercontent.com/leanprover/elan/master/elan-init.sh -sSf | bash -s -- -y
 source $HOME/.elan/env
diff --git a/scripts/build_example.sh b/scripts/build_example.sh
index db0577f..a194c38 100644
--- a/scripts/build_example.sh
+++ b/scripts/build_example.sh
@@ -4,9 +4,6 @@
 # 1. Launch a codespace for LeanInfer.
 # 2. Run `source scripts/build_example.sh`.
 
-# Install OpenBLAS.
-sudo apt-get update && sudo apt-get install -y libopenblas-dev
-
 # Set up elan.
 curl https://raw.githubusercontent.com/leanprover/elan/master/elan-init.sh -sSf | bash -s -- -y
 source $HOME/.elan/env

From 36b0aebd36602a080ee49d61b5f9df3730ebc7f1 Mon Sep 17 00:00:00 2001
From: Kaiyu Yang <kaiyuy@caltech.edu>
Date: Tue, 28 Nov 2023 08:05:29 -0800
Subject: [PATCH 3/4] cloud release

---
 lakefile.lean | 20 +++++++++++++++++---
 1 file changed, 17 insertions(+), 3 deletions(-)

diff --git a/lakefile.lean b/lakefile.lean
index f28560f..3f2648f 100644
--- a/lakefile.lean
+++ b/lakefile.lean
@@ -56,9 +56,13 @@ def hasCUDA : IO Bool := do
   return out.exitCode == 0
 
 
+def useCUDA : IO Bool := do
+  return (get_config? noCUDA |>.isNone) ∧ (← hasCUDA)
+
+
 def buildArchiveName : String :=
   let arch := if run_io isArm! then "arm64" else "x86_64"
-  if run_io hasCUDA then
+  if run_io useCUDA then
     s!"{arch}-cuda"
   else
     arch
@@ -206,7 +210,7 @@ target libopenblas pkg : FilePath := do
   afterReleaseAsync pkg do
     let rootDir := pkg.buildDir / "OpenBLAS"
     ensureDirExists rootDir
-    let dst := rootDir / (nameToSharedLib "openblas")
+    let dst := pkg.nativeLibDir / (nameToSharedLib "openblas")
     let url := "https://github.com/OpenMathLib/OpenBLAS"
 
     try
@@ -223,6 +227,16 @@ target libopenblas pkg : FilePath := do
           args := flags
           cwd := rootDir
         }
+        proc {
+          cmd := "cp"
+          args := #[(rootDir / nameToSharedLib "openblas").toString, dst.toString]
+        }
+        -- TODO: Don't hardcode the version "0".
+        let dst' := pkg.nativeLibDir / (nameToVersionedSharedLib "openblas" "0")
+        proc {
+          cmd := "cp"
+          args := #[dst.toString, dst'.toString]
+        }
       return (dst, trace)
 
     else
@@ -236,7 +250,7 @@ def getCt2CmakeFlags : IO (Array String) := do
   | .macos => flags := flags ++ #["-DWITH_ACCELERATE=ON", "-DWITH_OPENBLAS=OFF"]
   | .linux => flags := flags ++ #["-DWITH_ACCELERATE=OFF", "-DWITH_OPENBLAS=ON", "-DOPENBLAS_INCLUDE_DIR=../../OpenBLAS", "-DOPENBLAS_LIBRARY=../../OpenBLAS/libopenblas.so"]
 
-  if ← hasCUDA then
+  if ← useCUDA then
     flags := flags ++ #["-DWITH_CUDA=ON", "-DWITH_CUDNN=ON"]
   else
     flags := flags ++ #["-DWITH_CUDA=OFF", "-DWITH_CUDNN=OFF"]

From 8885cd68c4e4df232c543d53b75fbde8c9681b74 Mon Sep 17 00:00:00 2001
From: Kaiyu Yang <kaiyuy@caltech.edu>
Date: Tue, 28 Nov 2023 08:09:45 -0800
Subject: [PATCH 4/4] update readme

---
 README.md     | 10 +---------
 lakefile.lean | 26 --------------------------
 2 files changed, 1 insertion(+), 35 deletions(-)

diff --git a/README.md b/README.md
index 48baba2..97c2501 100644
--- a/README.md
+++ b/README.md
@@ -9,17 +9,9 @@ LeanInfer provides tactic suggestions by running LLMs through Lean's foreign fun
 ## Requirements
 
 * Supported platforms: Linux and macOS (:warning: maybe also Windows WSL, but untested)
-* xcode-select --install
-
-
 * Git LFS
-* A C++17 compatible compiler, e.g., recent versions of GCC or Clang
-* CMake >= 3.7
-* At least one of [MKL](https://www.intel.com/content/www/us/en/developer/tools/oneapi/onemkl.html), [Accelerate](https://developer.apple.com/documentation/accelerate), [DNNL](https://github.com/oneapi-src/oneDNN), or [OpenBLAS](https://www.openblas.net/). Most systems have them pre-installed, so you may not need to do anything. However, if you see the error `No BLAS library found`, it means you don't have any of these libraries. In that case, we recommend MKL for Intel CPUs, Accelerate for Apple, and DNNL for AMD CPUs. OpenBLAS is cross-platform and can be used as a last resort if you cannot install the others. We are happy to help if you have trouble installing these libraries.
 * Optional (recommended if you have a [CUDA-enabled GPU](https://developer.nvidia.com/cuda-gpus)): CUDA and [cuDNN](https://developer.nvidia.com/cudnn)
 
-**Please run `lake script run LeanInfer/check` to check if the requirements have been satisfied.**
-
 
 ## Adding LeanInfer as a Dependency to Your Project
 
@@ -27,7 +19,7 @@ LeanInfer provides tactic suggestions by running LLMs through Lean's foreign fun
 
 1. Add the package configuration option `moreLinkArgs := #["-L./.lake/packages/LeanInfer/.lake/build/lib", "-lonnxruntime", "-lctranslate2"]` to lakefile.lean. Also add LeanInfer as a dependency:
 ```lean
-require LeanInfer from git "https://github.com/lean-dojo/LeanInfer.git" @ "main"
+require LeanInfer from git "https://github.com/lean-dojo/LeanInfer.git" @ "v0.0.9"
 ```
 2. Run `lake update LeanInfer`
 3. Run `lake script run LeanInfer/download` to download the models from Hugging Face to `~/.cache/lean_infer/`
diff --git a/lakefile.lean b/lakefile.lean
index 3f2648f..203af08 100644
--- a/lakefile.lean
+++ b/lakefile.lean
@@ -443,31 +443,5 @@ script download do
   return 0
 
 
-def checkGitLFS : IO Bool := do
-  if ← checkAvailable "git" then
-    let proc ← IO.Process.output {
-      cmd := "git"
-      args := #["lfs", "install"]
-    }
-    return proc.exitCode == 0
-  else
-    return false
-
-
-script check do
-  if Platform.isWindows then
-    error "Windows is not supported"
-  if ¬ (← checkGitLFS) then
-    error "Git LFS is not installed"
-  if ¬ (← checkAvailable "c++") then
-    error "C++ compiler is not installed"
-  if ¬ (← checkAvailable "cmake") then
-    error "CMake is not installed"
-  if ¬ (← checkAvailable "make") then
-    error "Make is not installed"
-  println! "Looks good to me!"
-  return 0
-
-
 require std from git "https://github.com/leanprover/std4" @ "main"
 require aesop from git "https://github.com/JLimperg/aesop" @ "master"