File size: 1,153 Bytes
edbcf22 |
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 |
"""Load Hydra BitNet model."""
import torch
from safetensors.torch import load_file
def load_hydra(model_path: str, device: str = "cpu"):
"""Load Hydra model from HuggingFace format."""
import sys
from pathlib import Path
# Add aisim to path if needed
aisim_path = Path(__file__).parent.parent / "aisim"
if aisim_path.exists():
sys.path.insert(0, str(aisim_path))
from bitnet_moe import M2MSentinel
import json
# Load config
with open(f"{model_path}/config.json") as f:
config = json.load(f)
# Create model
model = M2MSentinel(
vocab_size=config["vocab_size"],
dim=config["hidden_size"],
depth=config["num_hidden_layers"],
experts=config["num_experts"],
)
# Load weights
weights = load_file(f"{model_path}/model.safetensors")
model.load_state_dict(weights)
model = model.to(device)
model.eval()
return model, config
if __name__ == "__main__":
import sys
model_path = sys.argv[1] if len(sys.argv) > 1 else "."
model, config = load_hydra(model_path)
print(f"Loaded model: {config}")
|