From 33fbac0880bab2f24649b1b3f533d0a95cc8bc39 Mon Sep 17 00:00:00 2001
From: Jan Wassenberg <janwas@google.com>
Date: Mon, 4 Aug 2025 22:35:59 -0700
Subject: [PATCH] Exporter updates/fixes

PiperOrigin-RevId: 791046073
---
 python/convert_from_safetensors.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/python/convert_from_safetensors.py b/python/convert_from_safetensors.py
index 6574c98..b67b5b8 100644
--- a/python/convert_from_safetensors.py
+++ b/python/convert_from_safetensors.py
@@ -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: