File size: 1,181 Bytes
ea416d9 | 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 | #!/usr/bin/env python3
"""Minimal checkpoint loading example for Tianmu-Emb-Uni-8B adapter weights.
This repository releases trained adapter/audio-side weights. Base model weights
for Qwen3-VL-Embedding-8B and Qwen2.5-Omni-7B must be available separately.
"""
from pathlib import Path
import sys
import torch
from safetensors.torch import load_file
def main():
repo_dir = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(repo_dir))
from tianmu_model.modeling import OmniEmbedModel
weight_path = repo_dir / "model.safetensors"
model = OmniEmbedModel(
audio_encoder_type="omni",
audio_model_path="/path/to/Qwen2.5-Omni-7B",
vl_model_name="/path/to/Qwen3-VL-Embedding-8B",
freeze_vl=True,
freeze_audio_encoder=True,
)
state_dict = load_file(str(weight_path), device="cpu")
missing, unexpected = model.load_state_dict(state_dict, strict=False)
print(f"loaded tensors: {len(state_dict)}")
print(f"missing keys: {len(missing)}")
print(f"unexpected keys: {len(unexpected)}")
model.eval()
with torch.no_grad():
print("model ready")
if __name__ == "__main__":
main()
|