()
| 548 | |
| 549 | |
| 550 | def main(): |
| 551 | parser = get_argument_parser() |
| 552 | args = parser.parse_args() |
| 553 | args_sanity_checks(args) |
| 554 | if args.dataset is None: |
| 555 | exp_name = f"{get_exp_name(args.config)}_{args.precision}_dummy_cal_{args.num_chunks}_chunks" |
| 556 | else: |
| 557 | exp_name = ( |
| 558 | f"{get_exp_name(args.config)}_{args.precision}_{args.num_chunks}_chunks" |
| 559 | ) |
| 560 | if args.platform == "DX4": |
| 561 | platform_b = b"mt6991" |
| 562 | elif args.platform == "DX3": |
| 563 | platform_b = b"mt6989" |
| 564 | else: |
| 565 | raise ValueError( |
| 566 | f"Platform should be either DX3 or DX4, but got {args.platform}" |
| 567 | ) |
| 568 | print_args(args, exp_name) |
| 569 | enc_exp_name = f"{get_exp_name(args.config)}_{args.precision}_encoder" |
| 570 | |
| 571 | config, weight_dir, tokenizer_class, chunk_class = resolve_model_classes( |
| 572 | args.config |
| 573 | ) |
| 574 | encoder_class = chunk_class[0] |
| 575 | decoder_class = chunk_class[1] |
| 576 | tokenizer = tokenizer_class.from_pretrained( |
| 577 | weight_dir, language="english", task="transcribe" |
| 578 | ) |
| 579 | |
| 580 | preprocessor_attr = config.p.__dict__ |
| 581 | preprocessor = WhisperAudioProcessor(**preprocessor_attr) |
| 582 | |
| 583 | head_dim = config.llm.head_dim |
| 584 | |
| 585 | # Evenly distribute the layers across chunks. |
| 586 | num_blocks_per_chunk = [ |
| 587 | (config.llm.num_hidden_layers // args.num_chunks) |
| 588 | + (i < (config.llm.num_hidden_layers % args.num_chunks)) |
| 589 | for i in range(args.num_chunks) |
| 590 | ] |
| 591 | check_all_chunks_same_num_layer(num_blocks_per_chunk) # noqa: F405 |
| 592 | |
| 593 | output_folder = os.path.join("pte", exp_name) |
| 594 | enc_output_folder = os.path.join("pte", enc_exp_name) |
| 595 | |
| 596 | # Load all collected checkpoint files into one giant state_dict |
| 597 | state_dict = load_checkpoints(weight_dir) |
| 598 | |
| 599 | dump_embedding_lut_for_cmdline(weight_dir, state_dict, config.llm) |
| 600 | |
| 601 | export_shapes, max_num_token, max_cache_size = get_export_shapes(args.shapes) |
| 602 | print(f"export shapes: {export_shapes}") |
| 603 | print(f"Max Num Token: {max_num_token}") |
| 604 | print(f"Max Cache Size: {max_cache_size}") |
| 605 | |
| 606 | embedding_layer = get_embedding_layer(config.llm, weight_dir, state_dict) |
| 607 |
no test coverage detected