Files
AstrAI/astrai/serialization/__init__.py
T
ViperEkura 7d27f3e078 feat: load HuggingFace checkpoints via key/config conversion
- Add astrai.serialization.hf_adapter mapping LLaMA-style HF keys to AstrAI names (input_layernorm, gate_proj, MoE experts/shared_experts) with config aliases for dense and MoE (Mixtral/DeepSeek-V3) layouts; reject biased projections, mismatched head_dim and MLA
- Give AutoModel.from_pretrained weights_format=auto|astrai|hf with auto-detection; read sharded safetensors via model.safetensors.index.json
- Adapt preloaded weights/config in train_context and benchmark CLI
2026-08-20 11:34:59 +08:00

54 lines
1.1 KiB
Python

"""Serialization utilities for models and datasets.
This package re-exports checkpoint helpers and dataset storage helpers so
that existing imports from ``astrai.serialization`` continue to work.
"""
from astrai.serialization.checkpoint import (
Checkpoint,
load_json,
load_model_config,
load_model_weights,
load_safetensors,
load_state_dict,
load_torch,
save_json,
save_model,
save_safetensors,
save_torch,
)
from astrai.serialization.dataset import (
load_bin,
load_bin_offsets,
save_bin,
)
from astrai.serialization.hf_adapter import (
HF_MODEL_TYPES,
adapt_config,
convert_hf_config,
convert_hf_weights,
looks_like_hf_state_dict,
)
__all__ = [
"Checkpoint",
"HF_MODEL_TYPES",
"adapt_config",
"convert_hf_config",
"convert_hf_weights",
"looks_like_hf_state_dict",
"load_json",
"load_model_config",
"load_model_weights",
"load_safetensors",
"load_state_dict",
"load_torch",
"save_json",
"save_model",
"save_safetensors",
"save_torch",
"load_bin",
"load_bin_offsets",
"save_bin",
]