(model_dir: Path, output_dir: Path)
| 76 | |
| 77 | |
| 78 | def extract(model_dir: Path, output_dir: Path) -> None: |
| 79 | output_dir.mkdir(parents=True, exist_ok=True) |
| 80 | backbone_dir = output_dir / "qwen3_backbone" |
| 81 | backbone_dir.mkdir(exist_ok=True) |
| 82 | embed_dir = output_dir / "embeddings" |
| 83 | embed_dir.mkdir(exist_ok=True) |
| 84 | head_dir = output_dir / "lm_heads" |
| 85 | head_dir.mkdir(exist_ok=True) |
| 86 | |
| 87 | moss_config = load_source_config(model_dir) |
| 88 | index = load_source_index(model_dir) |
| 89 | weight_map = index["weight_map"] |
| 90 | |
| 91 | lang_config = moss_config["language_config"] |
| 92 | n_vq = moss_config.get("n_vq", 32) |
| 93 | hidden_size = lang_config["hidden_size"] |
| 94 | vocab_size = lang_config["vocab_size"] |
| 95 | audio_vocab_size = moss_config.get("audio_vocab_size", 1024) |
| 96 | |
| 97 | log.info( |
| 98 | "Model: hidden_size=%d, vocab_size=%d, n_vq=%d, audio_vocab_size=%d", |
| 99 | hidden_size, vocab_size, n_vq, audio_vocab_size, |
| 100 | ) |
| 101 | |
| 102 | shard_to_tensors: dict[str, list[str]] = defaultdict(list) |
| 103 | for tensor_name, shard_file in weight_map.items(): |
| 104 | shard_to_tensors[shard_file].append(tensor_name) |
| 105 | |
| 106 | backbone_tensors: dict[str, torch.Tensor] = {} |
| 107 | backbone_size = 0 |
| 108 | shard_idx = 0 |
| 109 | saved_shards: list[str] = [] |
| 110 | backbone_weight_map: dict[str, str] = {} |
| 111 | |
| 112 | def flush_backbone_shard(): |
| 113 | nonlocal backbone_tensors, backbone_size, shard_idx |
| 114 | if not backbone_tensors: |
| 115 | return |
| 116 | shard_idx += 1 |
| 117 | shard_name = f"model-{shard_idx:05d}-of-PLACEHOLDER.safetensors" |
| 118 | shard_path = backbone_dir / shard_name |
| 119 | log.info(" Writing backbone shard %s (%d tensors, %.2f GB)", |
| 120 | shard_name, len(backbone_tensors), backbone_size / 1e9) |
| 121 | save_file(backbone_tensors, str(shard_path)) |
| 122 | for tname in backbone_tensors: |
| 123 | backbone_weight_map[tname] = shard_name |
| 124 | saved_shards.append(shard_name) |
| 125 | backbone_tensors = {} |
| 126 | backbone_size = 0 |
| 127 | |
| 128 | sorted_shards = sorted(shard_to_tensors.keys()) |
| 129 | for shard_file in sorted_shards: |
| 130 | tensor_names = shard_to_tensors[shard_file] |
| 131 | shard_path = model_dir / shard_file |
| 132 | log.info("Processing shard: %s (%d tensors)", shard_file, len(tensor_names)) |
| 133 | |
| 134 | with safe_open(str(shard_path), framework="pt") as sf: |
| 135 | for tname in sorted(tensor_names): |
no test coverage detected