File size: 4,987 Bytes
3f14c20 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 | """
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!")
|