Upload folder using huggingface_hub

config.json

@@ -6,6 +6,7 @@
     "auto_map": {
         "AutoConfig": "configuration_nanogpt.NanoGPTConfig",
         "AutoModel": "modeling_nanogpt.NanoGPTModel",
+        "AutoModelForCausalLM": "modeling_nanogpt.NanoGPTModel",
         "AutoTokenizer": "tokenizer_nanogpt.NanoGPTTokenizer"
     },
     "sequence_len": 2048,

tokenizer_nanogpt.py

@@ -1,5 +1,7 @@
 import os
 import pickle
+from huggingface_hub import hf_hub_download
+from huggingface_hub.utils import HfHubHTTPError
 
 
 class NanoGPTTokenizer:
@@ -20,9 +22,36 @@ class NanoGPTTokenizer:
 
     @classmethod
     def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
-        tok_path = os.path.join(pretrained_model_name_or_path, "tokenizer.pkl")
-        with open(tok_path, "rb") as f:
-            enc = pickle.load(f)
+        """
+        Load tokenizer from either:
+        - Local directory path
+        - Hugging Face Hub repo ID
+        - Cached directory (handled automatically)
+        """
+        # First, try to load from local path
+        local_tok_path = os.path.join(pretrained_model_name_or_path, "tokenizer.pkl")
+
+        if os.path.isfile(local_tok_path):
+            # Local file exists, load it directly
+            with open(local_tok_path, "rb") as f:
+                enc = pickle.load(f)
+        else:
+            # Try to download from Hugging Face Hub
+            try:
+                # This handles cache automatically and returns the cached file path
+                tok_path = hf_hub_download(
+                    repo_id=pretrained_model_name_or_path,
+                    filename="tokenizer.pkl",
+                    **kwargs
+                )
+                with open(tok_path, "rb") as f:
+                    enc = pickle.load(f)
+            except (HfHubHTTPError, OSError) as e:
+                raise ValueError(
+                    f"Could not load tokenizer.pkl from {pretrained_model_name_or_path}. "
+                    f"Make sure the path exists or the repo is accessible on the Hub."
+                ) from e
+
         return cls(enc)
 
     def encode(self, text, prepend=None):