extract.py

import torch
from safetensors.torch import load_file
import json
import os

# path to HF checkpoint
index_file = "model.safetensors.index.json"

with open(index_file, "r") as f:
    index = json.load(f)

# all shards
shards = index["weight_map"].values()
shards = sorted(list(set(shards)))

# load each shard
state_dict = {}
for shard in shards:
    print(f"Loading {shard}...")
    sd = load_file(shard)
    state_dict.update(sd)

# extract only the vision encoder keys
visual_sd = {
    k: v for k, v in state_dict.items()
    if k.startswith("visual.") or k.startswith("vision_tower.")
}

print(f"Extracted {len(visual_sd)} visual weights")

torch.save(visual_sd, "qwen_visual_model_weights.pth")

print("Saved to qwen_visual_model_weights.pth")