216 lines
7.1 KiB
Python
216 lines
7.1 KiB
Python
#!/usr/bin/env python3
|
|
import argparse
|
|
import json
|
|
import shutil
|
|
import struct
|
|
from pathlib import Path
|
|
|
|
|
|
def load_json(path: Path):
|
|
with open(path, "r", encoding="utf-8") as handle:
|
|
return json.load(handle)
|
|
|
|
|
|
def save_json(path: Path, value):
|
|
with open(path, "w", encoding="utf-8") as handle:
|
|
json.dump(value, handle, indent=2, ensure_ascii=False)
|
|
handle.write("\n")
|
|
|
|
|
|
def load_safetensors_header(path: Path):
|
|
with open(path, "rb") as handle:
|
|
header_len = struct.unpack("<Q", handle.read(8))[0]
|
|
header = json.loads(handle.read(header_len).decode("utf-8"))
|
|
return header_len, header
|
|
|
|
|
|
def tensor_payload_size(entry: dict) -> int:
|
|
start, end = entry["data_offsets"]
|
|
return end - start
|
|
|
|
|
|
def tensor_element_count(entry: dict) -> int:
|
|
count = 1
|
|
for dim in entry["shape"]:
|
|
count *= dim
|
|
return count
|
|
|
|
|
|
def build_vision_shard(
|
|
original_model_dir: Path,
|
|
original_index: dict,
|
|
output_path: Path,
|
|
):
|
|
weight_map = original_index["weight_map"]
|
|
visual_keys = sorted(key for key in weight_map if key.startswith("model.visual"))
|
|
if not visual_keys:
|
|
raise ValueError("original model has no model.visual tensors")
|
|
|
|
source_headers = {}
|
|
output_header = {}
|
|
output_payload = bytearray()
|
|
parameter_count = 0
|
|
|
|
for key in visual_keys:
|
|
shard_name = weight_map[key]
|
|
shard_path = original_model_dir / shard_name
|
|
if shard_path not in source_headers:
|
|
source_headers[shard_path] = load_safetensors_header(shard_path)
|
|
_, source_header = source_headers[shard_path]
|
|
entry = source_header[key]
|
|
data_start, data_end = entry["data_offsets"]
|
|
renamed_key = key.replace("model.visual", "vision_tower", 1)
|
|
parameter_count += tensor_element_count(entry)
|
|
|
|
with open(shard_path, "rb") as handle:
|
|
payload_base = 8 + source_headers[shard_path][0]
|
|
handle.seek(payload_base + data_start)
|
|
tensor_bytes = handle.read(data_end - data_start)
|
|
|
|
new_start = len(output_payload)
|
|
output_payload.extend(tensor_bytes)
|
|
new_end = len(output_payload)
|
|
output_header[renamed_key] = {
|
|
"dtype": entry["dtype"],
|
|
"shape": entry["shape"],
|
|
"data_offsets": [new_start, new_end],
|
|
}
|
|
|
|
header_bytes = json.dumps(output_header, separators=(",", ":"), ensure_ascii=False).encode(
|
|
"utf-8"
|
|
)
|
|
with open(output_path, "wb") as handle:
|
|
handle.write(struct.pack("<Q", len(header_bytes)))
|
|
handle.write(header_bytes)
|
|
handle.write(output_payload)
|
|
|
|
return {
|
|
"tensor_count": len(output_header),
|
|
"payload_size": len(output_payload),
|
|
"parameter_count": parameter_count,
|
|
"weight_map": {name: output_path.name for name in output_header},
|
|
}
|
|
|
|
|
|
def copy_support_files(original_model_dir: Path, output_dir: Path):
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
for item in original_model_dir.iterdir():
|
|
if item.name == "model.safetensors.index.json":
|
|
continue
|
|
if item.name.startswith("model-") and item.suffix == ".safetensors":
|
|
continue
|
|
destination = output_dir / item.name
|
|
if item.is_dir():
|
|
if destination.exists():
|
|
shutil.rmtree(destination)
|
|
shutil.copytree(item, destination)
|
|
else:
|
|
shutil.copy2(item, destination)
|
|
|
|
|
|
def copy_quantized_text_shards(quantized_model_dir: Path, output_dir: Path):
|
|
index = load_json(quantized_model_dir / "model.safetensors.index.json")
|
|
for shard_name in sorted(set(index["weight_map"].values())):
|
|
shutil.copy2(quantized_model_dir / shard_name, output_dir / shard_name)
|
|
return index
|
|
|
|
|
|
def merged_config(original_config: dict, quantized_config: dict):
|
|
config = dict(original_config)
|
|
config["text_config"] = quantized_config["text_config"]
|
|
config["quantization"] = quantized_config["quantization"]
|
|
config["quantization_config"] = quantized_config.get(
|
|
"quantization_config", quantized_config["quantization"]
|
|
)
|
|
if "dtype" in quantized_config:
|
|
config["dtype"] = quantized_config["dtype"]
|
|
return config
|
|
|
|
|
|
def build_merged_index(
|
|
quantized_index: dict,
|
|
vision_weight_map: dict,
|
|
vision_payload_size: int,
|
|
vision_parameter_count: int,
|
|
):
|
|
merged_weight_map = dict(quantized_index["weight_map"])
|
|
merged_weight_map.update(vision_weight_map)
|
|
|
|
metadata = dict(quantized_index.get("metadata", {}))
|
|
total_size = metadata.get("total_size")
|
|
if isinstance(total_size, str):
|
|
try:
|
|
total_size = int(float(total_size))
|
|
except ValueError:
|
|
total_size = None
|
|
if not isinstance(total_size, int):
|
|
total_size = 0
|
|
metadata["total_size"] = total_size + vision_payload_size
|
|
|
|
total_parameters = metadata.get("total_parameters")
|
|
if isinstance(total_parameters, str):
|
|
try:
|
|
total_parameters = int(float(total_parameters))
|
|
except ValueError:
|
|
total_parameters = None
|
|
if isinstance(total_parameters, int):
|
|
metadata["total_parameters"] = total_parameters + vision_parameter_count
|
|
|
|
return {
|
|
"metadata": metadata,
|
|
"weight_map": merged_weight_map,
|
|
}
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description="Build an MLX multimodal Qwen3.5-MoE model by combining a text-quantized MLX export with the original vision tower."
|
|
)
|
|
parser.add_argument("original_model_dir", type=Path)
|
|
parser.add_argument("quantized_text_dir", type=Path)
|
|
parser.add_argument("output_dir", type=Path)
|
|
parser.add_argument(
|
|
"--vision-shard-name",
|
|
default="vision-00001-of-00001.safetensors",
|
|
help="Filename for the rewritten vision safetensors shard.",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
original_model_dir = args.original_model_dir.resolve()
|
|
quantized_text_dir = args.quantized_text_dir.resolve()
|
|
output_dir = args.output_dir.resolve()
|
|
|
|
copy_support_files(original_model_dir, output_dir)
|
|
quantized_index = copy_quantized_text_shards(quantized_text_dir, output_dir)
|
|
|
|
original_index = load_json(original_model_dir / "model.safetensors.index.json")
|
|
vision_shard = build_vision_shard(
|
|
original_model_dir,
|
|
original_index,
|
|
output_dir / args.vision_shard_name,
|
|
)
|
|
|
|
quantized_config = load_json(quantized_text_dir / "config.json")
|
|
original_config = load_json(original_model_dir / "config.json")
|
|
save_json(output_dir / "config.json", merged_config(original_config, quantized_config))
|
|
save_json(
|
|
output_dir / "model.safetensors.index.json",
|
|
build_merged_index(
|
|
quantized_index,
|
|
vision_shard["weight_map"],
|
|
vision_shard["payload_size"],
|
|
vision_shard["parameter_count"],
|
|
),
|
|
)
|
|
|
|
print(f"wrote {output_dir}")
|
|
print(f"quantized_text_tensors={len(quantized_index['weight_map'])}")
|
|
print(f"vision_tensors={vision_shard['tensor_count']}")
|
|
print(f"vision_payload_bytes={vision_shard['payload_size']}")
|
|
print(
|
|
f"merged_tensors={len(quantized_index['weight_map']) + vision_shard['tensor_count']}"
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|