Upload folder using huggingface_hub

modeling_nanogpt.py

@@ -1,5 +1,7 @@
 import math
+import os
 from dataclasses import dataclass
+from pathlib import Path
 
 import torch
 import torch.nn as nn
@@ -170,5 +172,61 @@ class NanoGPTModel(PreTrainedModel):
             loss = F.cross_entropy(logits.view(-1, logits.size(-1)), labels.view(-1), ignore_index=-1, reduction='mean')
         return {"loss": loss, "logits": logits}
 
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
+        config = kwargs.pop("config", None)
+        subfolder = kwargs.pop("subfolder", None)
+        device_map = kwargs.get("device_map")
+        if device_map is not None:
+            # Delegate complex dispatch (like accelerate) to the base implementation.
+            if subfolder is not None:
+                kwargs["subfolder"] = subfolder
+            if config is not None:
+                kwargs["config"] = config
+            return super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
+
+        base_path = Path(pretrained_model_name_or_path)
+        if subfolder:
+            base_path = base_path / subfolder
+
+        weight_path = None
+        if base_path.is_dir():
+            candidate_files = [
+                base_path / "pytorch_model.bin",
+                base_path / "model.bin",
+            ]
+            candidate_files.extend(sorted(base_path.glob("model_*.pt"), reverse=True))
+            candidate_files.extend(sorted(base_path.glob("*.bin"), reverse=True))
+            for cand in candidate_files:
+                if cand.is_file():
+                    weight_path = cand
+                    break
+
+        if weight_path is None:
+            # Fall back to the default behaviour (e.g. remote repo or standard filenames)
+            if subfolder is not None:
+                kwargs["subfolder"] = subfolder
+            if config is not None:
+                kwargs["config"] = config
+            return super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
+
+        if config is None:
+            config = NanoGPTConfig.from_pretrained(pretrained_model_name_or_path, subfolder=subfolder)
+
+        torch_dtype = kwargs.pop("torch_dtype", None)
+        strict = kwargs.pop("strict", True)
+
+        state_dict = torch.load(str(weight_path), map_location="cpu")
+        if isinstance(state_dict, dict) and "state_dict" in state_dict:
+            state_dict = state_dict["state_dict"]
+        state_dict = {k.lstrip("_orig_mod."): v for k, v in state_dict.items()}
+
+        model = cls(config, *model_args)
+        model.load_state_dict(state_dict, strict=strict)
+        if torch_dtype is not None:
+            model = model.to(dtype=torch_dtype)
+        model.eval()
+        return model
+