()
| 831 | |
| 832 | |
| 833 | def main() -> None: |
| 834 | args = parse_args() |
| 835 | args = validate_args(args) |
| 836 | model_path = _require_namespace_path(args, "model_path") |
| 837 | codec_model_path = _require_namespace_path(args, "codec_model_path") |
| 838 | output_dir = _require_namespace_path(args, "save_path") |
| 839 | overwrite = _require_namespace_bool(args, "overwrite") |
| 840 | |
| 841 | main_config = load_json(model_path / "config.json") |
| 842 | codec_config = load_json(codec_model_path / "config.json") |
| 843 | main_index = load_source_index(model_path) |
| 844 | codec_index = load_source_index(codec_model_path) |
| 845 | tokenizer_filenames = collect_tokenizer_asset_filenames(model_path) |
| 846 | merged_config = build_merged_config(main_config, codec_config) |
| 847 | |
| 848 | if overwrite and output_dir.exists(): |
| 849 | cleanup_partial_output_dir(output_dir) |
| 850 | |
| 851 | output_dir.mkdir(parents=True, exist_ok=False) |
| 852 | |
| 853 | copied_tokenizer_files = copy_tokenizer_assets( |
| 854 | model_path, |
| 855 | output_dir, |
| 856 | tokenizer_filenames, |
| 857 | ) |
| 858 | copy_processor_source(output_dir) |
| 859 | write_processor_config(output_dir, merged_config) |
| 860 | patch_tokenizer_config_processor_class( |
| 861 | output_dir / TOKENIZER_CONFIG_FILENAME, |
| 862 | PROCESSOR_CLASS_NAME, |
| 863 | ) |
| 864 | |
| 865 | write_merged_config(output_dir, merged_config) |
| 866 | |
| 867 | main_weight_map = require_weight_map(main_index, "main_index") |
| 868 | remapped_main_weight_map = build_remapped_main_weight_map(main_weight_map) |
| 869 | codec_weight_map = require_weight_map(codec_index, "codec_index") |
| 870 | main_groups = group_weight_map_by_shard(main_weight_map) |
| 871 | codec_groups = group_weight_map_by_shard(codec_weight_map) |
| 872 | |
| 873 | expected_codec_weight_map = build_prefixed_codec_keys( |
| 874 | codec_weight_map, |
| 875 | remapped_main_weight_map, |
| 876 | CODEC_PREFIX, |
| 877 | ) |
| 878 | |
| 879 | merged_weight_map = write_merged_shards( |
| 880 | output_dir, |
| 881 | iter_merged_tensors( |
| 882 | model_path, |
| 883 | main_groups, |
| 884 | codec_model_path, |
| 885 | codec_groups, |
| 886 | CODEC_PREFIX, |
| 887 | ), |
| 888 | ) |
| 889 | |
| 890 | expected_merged_keys = set(remapped_main_weight_map) | set(expected_codec_weight_map) |
no test coverage detected