From e2e1931709d2b04a8f946b51c0f9b00f91ab9136 Mon Sep 17 00:00:00 2001
From: Rishabh Srivastava <rishabhsriv@gmail.com>
Date: Thu, 8 Feb 2024 21:54:27 +0800
Subject: [PATCH] updated launch CLI function to improve first time download
 experiences

---
 sqlcoder/cli.py | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

diff --git a/sqlcoder/cli.py b/sqlcoder/cli.py
index ea92b3d..692dc79 100644
--- a/sqlcoder/cli.py
+++ b/sqlcoder/cli.py
@@ -2,6 +2,7 @@ import os
 import sys
 import sqlcoder
 import subprocess
+from huggingface_hub import snapshot_download
 
 USAGE_STRING = """
 Usage: sqlcoder <command>
@@ -57,6 +58,24 @@ def serve_static():
         httpd.serve_forever()
 
 def launch():
+    home_dir = os.path.expanduser("~")
+    defog_path = os.path.join(home_dir, ".defog")
+    if not os.popen("lspci | grep -i nvidia").read():
+        # not a GPU machine
+        filepath = os.path.join(home_dir, ".defog", "sqlcoder-7b-q5_k_m.gguf")
+        if not os.path.exists(filepath):
+            print(
+                "Downloading the SQLCoder-7b-2 GGUF file. This is a ~5GB file and may take a long time to download. But once it's downloaded, it will be saved on your machine and you won't have to download it again."
+            )
+            sqlcoder.hf_hub_download(repo_id="defog/sqlcoder-7b-2", filename="sqlcoder-7b-q5_k_m.gguf", local_dir=defog_path)
+    else:
+        # check if the huggingface model is already downloaded from hub. If not, download it
+        from huggingface_hub import snapshot_download
+        print(
+            "Downloading the SQLCoder-7b-2 model. This is a ~14GB file and may take a long time to download. But once it's downloaded, it will be saved on your machine and you won't have to download it again."
+        )
+        _ = snapshot_download("defog/sqlcoder-7b-2")
+    
     print("Starting SQLCoder server...")
     static_process = subprocess.Popen(["sqlcoder", "serve-static"])