Update tools/download/models.py

tools/download/models.py

@@ -1,6 +1,6 @@
 from pathlib import Path
 
-from huggingface_hub import hf_hub_download
+from huggingface_hub import hf_hub_download, snapshot_download
 
 REPO_ID = "lucas-ventura/chapter-llama"
 
@@ -42,6 +42,40 @@ def download_model(model_id_or_path, overwrite=False, local_dir=None):
     return str(Path(cache_path).parent)
 
 
+def download_base_model(
+    repo_id="meta-llama/Meta-Llama-3.1-8B-Instruct",
+    local_dir="./models",
+    use_symlinks=False,
+):
+    """
+    Downloads the base model from Hugging Face Hub.
+
+    Args:
+        repo_id (str): The repository ID on Hugging Face
+        local_dir (str): Directory to save the model to
+        use_symlinks (bool): Whether to use symlinks for the downloaded files
+
+    Returns:
+        str: Path to the downloaded model directory
+    """
+    try:
+        print(f"Downloading {repo_id} to {local_dir}...")
+        model_path = snapshot_download(
+            repo_id=repo_id, local_dir=local_dir, local_dir_use_symlinks=use_symlinks
+        )
+        print(f"Model downloaded successfully to: {model_path}")
+        return model_path
+    except Exception as e:
+        print(f"Error downloading model {repo_id}: {e}")
+        print(
+            f"\nYou are downloading `{repo_id}` to `{local_dir}` but failed. "
+            f"Please accept the agreement and obtain access at https://huggingface.co/{repo_id}. "
+            f"Then, use `huggingface-cli login` and your access tokens at https://huggingface.co/settings/tokens to authenticate. "
+            f"After that, run the code again."
+        )
+        return None
+
+
 if __name__ == "__main__":
     import argparse
 
@@ -67,4 +101,4 @@ if __name__ == "__main__":
     model_dir = download_model(
         args.model_id, overwrite=args.overwrite, local_dir=args.local_dir
     )
-    print(f"Model directory: {model_dir}")
+    print(f"Model directory: {model_dir}")