-
Notifications
You must be signed in to change notification settings - Fork 2
/
start.py
37 lines (27 loc) · 1.17 KB
/
start.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
import os
import shlex
import subprocess
import time
from pathlib import Path
from huggingface_hub import snapshot_download
def setup_dir(base_path: str):
dir_list = ["var/lib/tmp/data", "var/lib/tmp/jobs", "var/lib/tmp/.cache", "models", "db/sqlite"]
for _dl in dir_list:
p = Path(f"{base_path}/{_dl}")
if not p.is_dir():
p.mkdir(parents=True, exist_ok=True)
print(f"Download models...")
base_path = (Path(__file__).parent).resolve() if os.path.isdir("./.sidekickvenv/bin/") else "/meta_data"
setup_dir(base_path)
# Model 1:
print(f"Download model 1...")
snapshot_download(repo_id="NumbersStation/nsql-llama-2-7B", cache_dir=f"{base_path}/models/")
# Model 2:
print(f"Download model 2...")
snapshot_download(repo_id="defog/sqlcoder-34b-alpha", cache_dir=f"{base_path}/models/")
print(f"Download embedding model...")
snapshot_download(repo_id="BAAI/bge-base-en", cache_dir=f"{base_path}/models/sentence_transformers/")
print("Starting SQL-Sidekick.")
DAEMON_PATH = "./.sidekickvenv/bin/uvicorn" if os.path.isdir("./.sidekickvenv/bin/") else "/resources/venv/bin/uvicorn"
cmd = f"{DAEMON_PATH} ui.app:main"
subprocess.check_output(shlex.split(cmd))