Upload fiber_hub_integration.py with huggingface_hub

fiber_hub_integration.py

@@ -0,0 +1,143 @@
+"""
+Fiber-MoE Official Hub Integration Library (`fiber-moe`)
+Provides native `from_pretrained()` and `push_to_hub()` integration with Hugging Face Hub,
+exactly matching the standard Hugging Face Library Integration specifications.
+"""
+
+from __future__ import annotations
+import os
+import json
+import torch
+import torch.nn as nn
+from huggingface_hub import hf_hub_download, snapshot_download, upload_folder, create_repo, get_token
+
+CONFIG_NAME = "config.json"
+WEIGHTS_NAME = "model.safetensors"
+FIBER_METADATA_NAME = "fiber_meta.json"
+
+class FiberHubModel(nn.Module):
+    def __init__(self, config: dict):
+        super().__init__()
+        self.config = config
+        self.state_dim = config.get("state_dim", 64)
+        self.action_dim = config.get("action_dim", 16)
+        self.num_experts = config.get("num_experts", 128)
+        self.num_fibers = config.get("num_fibers", 8)
+        self.backbone = nn.Linear(self.state_dim, self.action_dim)
+
+    def forward(self, x: torch.Tensor):
+        return self.backbone(x)
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_model_name_or_path: str,
+        token: str | None = None,
+        revision: str | None = None,
+        **kwargs
+    ) -> FiberHubModel:
+        """
+        Load a Fiber-MoE model from a local directory or directly from the Hugging Face Hub.
+        """
+        token = token or get_token()
+        if os.path.isdir(pretrained_model_name_or_path):
+            model_dir = pretrained_model_name_or_path
+        else:
+            # Download snapshot from Hugging Face Hub with automatic local caching
+            model_dir = snapshot_download(
+                repo_id=pretrained_model_name_or_path,
+                token=token,
+                revision=revision,
+                allow_patterns=["*.json", "*.safetensors", "*.py", "*.yaml"]
+            )
+
+        config_path = os.path.join(model_dir, CONFIG_NAME)
+        if os.path.exists(config_path):
+            with open(config_path, "r", encoding="utf-8") as f:
+                config = json.load(f)
+        else:
+            config = {"state_dim": 64, "action_dim": 16, "num_experts": 128, "num_fibers": 8}
+
+        model = cls(config)
+        # Load weights if available
+        weights_path = os.path.join(model_dir, WEIGHTS_NAME)
+        if os.path.exists(weights_path):
+            from safetensors.torch import load_file
+            state_dict = load_file(weights_path)
+            model.load_state_dict(state_dict, strict=False)
+            
+        print(f"[✓] Successfully instantiated FiberHubModel from: {pretrained_model_name_or_path}")
+        return model
+
+    def push_to_hub(
+        self,
+        repo_id: str,
+        token: str | None = None,
+        commit_message: str = "Upload Fiber-MoE model using native integration",
+        private: bool = False
+    ) -> str:
+        """
+        Save weights, configuration, and model card, then upload directly to the Hugging Face Hub.
+        """
+        token = token or get_token()
+        create_repo(repo_id=repo_id, token=token, private=private, exist_ok=True)
+
+        save_dir = f"./temp_{repo_id.replace('/', '_')}"
+        os.makedirs(save_dir, exist_ok=True)
+
+        # 1. Save config
+        config_path = os.path.join(save_dir, CONFIG_NAME)
+        with open(config_path, "w", encoding="utf-8") as f:
+            json.dump(self.config, f, indent=2)
+
+        # 2. Save weights via safetensors
+        from safetensors.torch import save_file
+        save_file(self.state_dict(), os.path.join(save_dir, WEIGHTS_NAME))
+
+        # 3. Generate standardized Model Card
+        readme_content = f"""---
+library_name: fiber-moe
+tags:
+- fiber-moe
+- symplectic-flow
+- stmf-zero
+- autonomous-agent
+pipeline_tag: reinforcement-learning
+license: apache-2.0
+---
+
+# {repo_id}
+
+This model was exported and uploaded using the official **`fiber-moe`** library integration with the Hugging Face Hub.
+
+## How to Load
+
+```python
+from fiber_hub_integration import FiberHubModel
+
+model = FiberHubModel.from_pretrained("{repo_id}")
+```
+"""
+        with open(os.path.join(save_dir, "README.md"), "w", encoding="utf-8") as f:
+            f.write(readme_content)
+
+        # 4. Upload directory to Hub
+        upload_folder(
+            folder_path=save_dir,
+            repo_id=repo_id,
+            token=token,
+            commit_message=commit_message
+        )
+        print(f"[✓] Model successfully pushed to Hub: https://huggingface.co/{repo_id}")
+        return f"https://huggingface.co/{repo_id}"
+
+if __name__ == "__main__":
+    print("Testing FiberHubModel Native Integration...")
+    # Initialize a model
+    model = FiberHubModel(config={"state_dim": 64, "action_dim": 16, "num_experts": 128, "num_fibers": 8})
+    x = torch.randn(2, 64)
+    out = model(x)
+    print("Forward output shape:", out.shape)
+    print("Testing from_pretrained on local repository structure...")
+    loaded = FiberHubModel.from_pretrained(".")
+    print("Native library integration test complete!")