Exporter updates/fixes

PiperOrigin-RevId: 791046073
This commit is contained in:
Jan Wassenberg 2025-08-04 22:35:59 -07:00 committed by Copybara-Service
parent 4e062d68f7
commit 33fbac0880
1 changed files with 3 additions and 2 deletions

View File

@ -330,7 +330,7 @@ def export_paligemma_sbs(
"language_model.model.embed_tokens.weight"
][:-64]
writer = compression.SbsWriter()
writer = compression.SbsWriter(sbs_file)
metadata = []
scales = {}
dims = _get_dimensions(params)
@ -511,7 +511,8 @@ def export_paligemma_sbs(
# Write everything to the sbs file.
assert model_specifier.startswith("paligemma")
writer.write(configs.ModelConfig(model_specifier), tokenizer_file, sbs_file)
sbs_config = configs.ModelConfig(model_specifier)
writer.write(sbs_config, tokenizer_file)
# Write the metadata for manual inspection.
with open(csv_file, "w") as csv_handle: