| |
| |
| """Refine the proven 630-entry candidate to the native 70-token grid bound.""" |
|
|
| import argparse |
| import hashlib |
| import json |
| from pathlib import Path |
|
|
| from ai_edge_litert import schema_py_generated as schema |
| from litert_lm_builder import LitertLmFileBuilder |
|
|
| from build_size_candidates import crop_positions, pack_model |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--candidate-dir", type=Path, required=True) |
| args = parser.parse_args() |
| unpacked = args.candidate_dir / "unpacked" |
| report_path = args.candidate_dir / "build-report.json" |
| report = json.loads(report_path.read_text(encoding="utf-8")) |
| assert not any(x["name"] == "position210" for x in report["candidates"]) |
| model_path, = unpacked.glob("*TFLiteModel*tf_lite_vision_encoder*.tflite") |
| weight_path, = unpacked.glob("*TFLiteWeights*tf_lite_vision_encoder*.weight") |
| builder = LitertLmFileBuilder.from_toml_file(str(unpacked / "model.toml")) |
| model = schema.ModelT.InitFromObj(schema.Model.GetRootAsModel(model_path.read_bytes(), 0)) |
| refinement = crop_positions(model, weight_path, 210) |
| refinement["native_bound_argument"] = { |
| "patch_size": 16, "pooling_kernel_size": 3, "max_subpatches": 630, |
| "preprocessed_dimensions_are_multiples_of_pixels": 48, |
| "minimum_patch_grid_side": 3, |
| "maximum_patch_grid_side": 210, |
| "valid_position_coordinate_range": [0, 209], |
| "source": "https://github.com/google-ai-edge/LiteRT-LM/blob/v0.18.0/support/preprocessor/image_preprocessor_utils.cc", |
| "scope": "Current app's native image preprocessing with vision_tokens_per_image=70 only.", |
| } |
| model_path.write_bytes(pack_model(model)) |
| output_path = args.candidate_dir / "eg2-gallery-position210.litertlm" |
| with output_path.open("wb") as destination: |
| builder.build(destination) |
| report["position210_refinement"] = refinement |
| candidate = {"name": "position210", "path": str(output_path), "bytes": output_path.stat().st_size, |
| "saved_bytes": report["source_bytes"] - output_path.stat().st_size, |
| "sha256": hashlib.file_digest(output_path.open("rb"), "sha256").hexdigest()} |
| report["candidates"].append(candidate) |
| report_path.write_text(json.dumps(report, indent=2), encoding="utf-8") |
| print(json.dumps(candidate, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
