| """Derive a CTC-only encoder by removing the unused TDT projection output. |
| |
| No weight values or precision settings change. CoreML compilation prunes the |
| unreferenced projection operations. Probability and length outputs stay intact. |
| """ |
| import argparse |
| import hashlib |
| import json |
| from pathlib import Path |
| import shutil |
|
|
| import coremltools as ct |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--source", type=Path, required=True) |
| parser.add_argument("--out", type=Path, required=True) |
| args = parser.parse_args() |
| assert not args.out.exists() |
| args.out.mkdir(parents=True) |
| source = ct.models.MLModel(str(args.source / "Encoder.mlpackage"), skip_model_load=True) |
| spec = source.get_spec() |
| descriptions = [spec.description, *spec.description.functions] |
| removed = [] |
| for description in descriptions: |
| for index in reversed(range(len(description.output))): |
| if description.output[index].name == "encoder_proj": |
| removed.append(getattr(description, "name", "default")) |
| del description.output[index] |
| for name, function in spec.mlProgram.functions.items(): |
| for specialization, block in function.block_specializations.items(): |
| outputs = [x for x in block.outputs if x != "encoder_proj"] |
| assert len(outputs) == 2, (name, specialization, list(block.outputs)) |
| del block.outputs[:] |
| block.outputs.extend(outputs) |
| assert removed |
| result = ct.models.MLModel(spec, weights_dir=source.weights_dir, skip_model_load=True) |
| result.save(str(args.out / "Encoder.mlpackage")) |
| shutil.copytree(args.source / "Preprocessor.mlpackage", args.out / "Preprocessor.mlpackage") |
| shutil.copyfile(args.source / "vocab.json", args.out / "vocab.json") |
| config = json.loads((args.source / "config.json").read_text()) |
| config.pop("tdt_decoder", None) |
| config["models"]["Encoder.mlmodelc"]["outputs"].pop("encoder_proj", None) |
| config["models"].pop("DecoderJoint.mlmodelc", None) |
| config["models"].pop("JointBlock.mlmodelc", None) |
| config["model"] = "parakeet-ctc-110m-coreml" |
| config["version"] = "1.1.0-ctc-candidate" |
| (args.out / "config.json").write_text(json.dumps(config, indent=2) + "\n") |
| before = Path(source.weights_dir) / "weight.bin" |
| after = next((args.out / "Encoder.mlpackage").rglob("weight.bin")) |
| assert hashlib.sha256(before.read_bytes()).digest() == hashlib.sha256(after.read_bytes()).digest() |
| (args.out / "derivation.json").write_text(json.dumps({"removed_output": "encoder_proj", "functions": sorted(spec.mlProgram.functions), |
| "encoder_weight_sha256": hashlib.sha256(after.read_bytes()).hexdigest(), "weight_bytes_unchanged": True}, indent=2) + "\n") |
| print("CTC-only encoder derived with unchanged weights:", args.out) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
