Add download script for tinyllamas

README.md

@@ -57,7 +57,7 @@ Then login with `huggingface-cli login`
 ## Downloading Weights
 Models tested/supported
 ```text
-tinyllamas/stories{15,42,100}
+tinyllamas/stories{15M,42M,100M}
 openlm-research/open_llama_7b
 meta-llama/Llama-2-7b-chat-hf
 meta-llama/Llama-2-13b-chat-hf

scripts/convert_hf_checkpoint.py

@@ -24,6 +24,9 @@ def convert_hf_checkpoint(
     checkpoint_dir: Path = Path("checkpoints/meta-Transformer/Transformer-2-7b-chat-hf"),
     model_name: Optional[str] = None,
 ) -> None:
+    if "stories" in checkpoint_dir.name:
+        # No need to convert tinyllamas
+        return
     if model_name is None:
         model_name = checkpoint_dir.name
 

scripts/download.py

@@ -7,11 +7,28 @@
 from typing import Optional
 
 from requests.exceptions import HTTPError
+import subprocess
 
+def download_tinyllamas(repo_id: str, local_dir: str) -> None:
+    try:
+        model_name = repo_id.split("/")[-1]
+        # Download model weight
+        weight_url = "https://huggingface.co/karpathy/tinyllamas/resolve/main/" + model_name + ".pt"
+        weight_dst_path = os.path.join(local_dir, "model.pth")
+        subprocess.run(["wget", weight_url, "-O", weight_dst_path], check=True)
+        # Download tokenizer model
+        tokenizer_url = "https://github.com/karpathy/llama2.c/blob/master/tokenizer.model"
+        tokenizer_dst_path: str = os.path.join(local_dir, "tokenizer.model")
+        subprocess.run(["wget", tokenizer_url, "-O", tokenizer_dst_path], check=True)
+    except subprocess.CalledProcessError as e:
+        print(f"Failed to download {repo_id}: {e}")
 
 def hf_download(repo_id: Optional[str] = None, hf_token: Optional[str] = None) -> None:
     from huggingface_hub import snapshot_download
     os.makedirs(f"checkpoints/{repo_id}", exist_ok=True)
+    if "stories" in repo_id:
+        download_tinyllamas(repo_id, f"checkpoints/{repo_id}")
+        return
     try:
         snapshot_download(repo_id, local_dir=f"checkpoints/{repo_id}", local_dir_use_symlinks=False, token=hf_token)
     except HTTPError as e: