File size: 5,375 Bytes
6013524 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 | # Copyright 2026 Spawn Lau
# SPDX-License-Identifier: Apache-2.0
# Release change: optional fixed container metadata for byte-exact rebuilding.
"""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):
# LiteRT-LM builder 0.18.0 normally generates a new UUID and timestamp on
# every pack. Pin these non-computational fields to the tested release.
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()
|