Generalize model download #44

Author: pramitchoudhary

start.py

@@ -1,19 +1,24 @@
 import os
 import shlex
 import subprocess
+import time
 from pathlib import Path
 
 from huggingface_hub import snapshot_download
 
 print(f"Download model...")
 base_path = (Path(__file__).parent).resolve()
 
-# Model 1:
-print("Downloading model 1...")
-snapshot_download(repo_id="NumbersStation/nsql-llama-2-7B", cache_dir=f"{base_path}/models/")
-# Model 2:
-print("Downloading model 2...")
-snapshot_download(repo_id="defog/sqlcoder2", cache_dir=f"{base_path}/models/")
+MODEL_CHOICE_MAP = {
+    "h2ogpt-sql-sqlcoder2": "defog/sqlcoder2",
+    "h2ogpt-sql-nsql-llama-2-7B": "NumbersStation/nsql-llama-2-7B",
+}
+
+for _m in MODEL_CHOICE_MAP.values():
+    print(f"Downloading {_m}...", flush=True)
+    snapshot_download(repo_id=_m, cache_dir=f"{base_path}/models/")
+    time.sleep(3)
+
 print(f"Download embedding model...")
 snapshot_download(repo_id="BAAI/bge-base-en", cache_dir=f"{base_path}/models/sentence_transformers/")