| |
| |
| |
| """Remove unreferenced buffer records and deduplicate identical inline data.""" |
|
|
| import argparse |
| import hashlib |
| import json |
| from pathlib import Path |
|
|
| import flatbuffers |
| import numpy as np |
| from ai_edge_litert import schema_py_generated as schema |
| from litert_lm_builder import DType, LitertLmFileBuilder, Metadata |
| from litert_lm_builder import litertlm_builder as builder_impl |
|
|
|
|
| def published_system_metadata(entries): |
| |
| |
| return [entry for entry in entries if entry.key not in {"uuid", "creation_timestamp"}] + [ |
| Metadata(key="uuid", value="d8074fb5-3ec2-4550-9558-2c103b3ab73e", dtype=DType.STRING), |
| Metadata(key="creation_timestamp", value="2026-10-20T12:27:29.752155+00:00", dtype=DType.STRING), |
| ] |
|
|
|
|
| class SharedStringBuilder(flatbuffers.Builder): |
| def CreateString(self, value, encoding="utf-8", errors="strict"): |
| encoded = value.encode(encoding, errors) if isinstance(value, str) else bytes(value) |
| if self.sharedStrings is None: |
| self.sharedStrings = {} |
| if encoded not in self.sharedStrings: |
| self.sharedStrings[encoded] = super().CreateString(encoded) |
| return self.sharedStrings[encoded] |
|
|
|
|
| def pack(model): |
| builder = SharedStringBuilder(1024) |
| builder.Finish(model.Pack(builder), file_identifier=b"TFL3") |
| return bytes(builder.Output()) |
|
|
|
|
| def compact(path): |
| raw = path.read_bytes() |
| model = schema.ModelT.InitFromObj(schema.Model.GetRootAsModel(raw, 0)) |
| before_buffers = len(model.buffers) |
| before_external = len(model.externalBuffers or []) |
| used = {0} |
| used_external = set() |
| for graph in model.subgraphs: |
| for tensor in graph.tensors: |
| assert not tensor.isVariable |
| used.add(int(tensor.buffer)) |
| if tensor.externalBuffer: |
| used_external.add(int(tensor.externalBuffer)) |
| for metadata in model.metadata or []: |
| used.add(int(metadata.buffer)) |
| used.update(int(i) for i in model.metadataBuffer or []) |
| mapping, keys, buffers = {}, {}, [] |
| for old in sorted(used): |
| buffer = model.buffers[old] |
| key = (bytes(buffer.data) if buffer.data is not None else b"", int(buffer.offset), int(buffer.size)) |
| if key not in keys: |
| keys[key] = len(buffers) |
| buffers.append(buffer) |
| mapping[old] = keys[key] |
| for graph in model.subgraphs: |
| for tensor in graph.tensors: |
| old_buffer = model.buffers[int(tensor.buffer)] |
| new_buffer = buffers[mapping[int(tensor.buffer)]] |
| assert bytes(old_buffer.data or b"") == bytes(new_buffer.data or b"") if not isinstance(old_buffer.data, np.ndarray) else np.array_equal(old_buffer.data, new_buffer.data) |
| tensor.buffer = mapping[int(tensor.buffer)] |
| for metadata in model.metadata or []: |
| metadata.buffer = mapping[int(metadata.buffer)] |
| if model.metadataBuffer is not None: |
| model.metadataBuffer = np.asarray([mapping[int(i)] for i in model.metadataBuffer], dtype=np.int32) |
| model.buffers = buffers |
| model.externalBuffers = [e for e in model.externalBuffers or [] if int(e.id) in used_external] |
| result = pack(model) |
| path.write_bytes(result) |
| return {"file": path.name, "before_graph_bytes": len(raw), "after_graph_bytes": len(result), |
| "before_buffer_records": before_buffers, "after_buffer_records": len(buffers), |
| "before_external_buffer_records": before_external, |
| "after_external_buffer_records": len(model.externalBuffers), |
| "tensor_names_retained": True, "referenced_constant_values_retained": True} |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--candidate-dir", type=Path, required=True) |
| parser.add_argument("--reproducible", action="store_true", |
| help="Use the published container UUID and timestamp.") |
| args = parser.parse_args() |
| if args.reproducible: |
| builder_impl.populate_system_metadata = published_system_metadata |
| report_path = args.candidate_dir / "build-report.json" |
| report = json.loads(report_path.read_text(encoding="utf-8")) |
| assert not any(c["name"] == "compact_final" for c in report["candidates"]) |
| unpacked = args.candidate_dir / "unpacked" |
| builder = LitertLmFileBuilder.from_toml_file(str(unpacked / "model.toml")) |
| details = [compact(path) for path in sorted(unpacked.glob("*.tflite"))] |
| output = args.candidate_dir / "eg2-gallery-compact.litertlm" |
| with output.open("wb") as target: |
| builder.build(target) |
| candidate = {"name": "compact_final", "path": str(output), "bytes": output.stat().st_size, |
| "saved_bytes": report["source_bytes"] - output.stat().st_size, |
| "sha256": hashlib.file_digest(output.open("rb"), "sha256").hexdigest()} |
| report["metadata_compaction"] = details |
| report["candidates"].append(candidate) |
| report_path.write_text(json.dumps(report, indent=2), encoding="utf-8") |
| print(json.dumps({"candidate": candidate, "details": details}, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
