examples/load_weights.py

#!/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()