From ac608be800c2ba69acad5b6345e48a0c1d5065bf Mon Sep 17 00:00:00 2001
From: Roland Tannous <rolandtannous@gravityq.ai>
Date: Sat, 7 Mar 2026 10:57:51 +0000
Subject: [PATCH] fix: patch TokenizersBackend in export output after
 save_pretrained

---
 studio/backend/core/export/export.py | 38 ++++++++++++++++++++++++++++
 1 file changed, 38 insertions(+)

diff --git a/studio/backend/core/export/export.py b/studio/backend/core/export/export.py
index 9d9eff2548..6899eb949e 100644
--- a/studio/backend/core/export/export.py
+++ b/studio/backend/core/export/export.py
@@ -96,6 +96,7 @@ class ExportBackend:
         self.current_tokenizer = None
         self.is_vision = False
         self.is_peft = False
+        self._resolved_model_name = ""
 
     def cleanup_memory(self):
         """Offload and delete all models from memory"""
@@ -167,6 +168,7 @@ class ExportBackend:
 
             # Resolve model name for tokenizer patching (base model for LoRA, path otherwise)
             resolved_model_name = base_model or checkpoint_path
+            self._resolved_model_name = resolved_model_name
 
             # Patch broken tokenizer_config.json on disk before loading.
             # Qwen3.5/GLM checkpoints saved by TRL inherit "TokenizersBackend"
@@ -202,6 +204,23 @@ class ExportBackend:
             from utils.transformers_version import patch_tokenizer_in_memory
             patch_tokenizer_in_memory(tokenizer, model_name=resolved_model_name)
 
+            # Wrap tokenizer.save_pretrained so that every subsequent call
+            # (including internal ones from save_pretrained_merged /
+            # save_pretrained_gguf) auto-patches the on-disk output.
+            # Without this, the GGUF converter subprocess fails because
+            # save_pretrained re-writes "TokenizersBackend" to the output dir.
+            _orig_tok_save = tokenizer.save_pretrained
+            _fix_model_name = resolved_model_name
+
+            def _save_and_patch(*args, **kwargs):
+                result = _orig_tok_save(*args, **kwargs)
+                _dir = args[0] if args else kwargs.get("save_directory")
+                if _dir:
+                    patch_tokenizer_config(str(_dir), model_name=_fix_model_name)
+                return result
+
+            tokenizer.save_pretrained = _save_and_patch
+
             # Check if PEFT model
             self.is_peft = isinstance(model, (PeftModel, PeftModelForCausalLM))
 
@@ -279,6 +298,10 @@ class ExportBackend:
                     save_method=save_method
                 )
 
+                # Fix broken tokenizer_class on disk (belt-and-suspenders)
+                from utils.transformers_version import patch_tokenizer_config
+                patch_tokenizer_config(save_directory, model_name=self._resolved_model_name)
+
                 # Write export metadata so the Chat page can identify the base model
                 self._write_export_metadata(save_directory)
                 logger.info(f"Model saved successfully to {save_directory}")
@@ -335,6 +358,10 @@ class ExportBackend:
                 self.current_model.save_pretrained(save_directory)
                 self.current_tokenizer.save_pretrained(save_directory)
 
+                # Fix broken tokenizer_class on disk (belt-and-suspenders)
+                from utils.transformers_version import patch_tokenizer_config
+                patch_tokenizer_config(save_directory, model_name=self._resolved_model_name)
+
                 # Write export metadata so the Chat page can identify the base model
                 self._write_export_metadata(save_directory)
                 logger.info(f"Model saved successfully to {save_directory}")
@@ -446,6 +473,12 @@ class ExportBackend:
                     quantization_method=quant_method
                 )
 
+                # Fix broken tokenizer_class in intermediate HF output
+                # (save_pretrained wrapper handles pre-converter, this is
+                # belt-and-suspenders for the final on-disk state)
+                from utils.transformers_version import patch_tokenizer_config
+                patch_tokenizer_config(model_save_path, model_name=self._resolved_model_name)
+
                 # Relocate GGUF artifacts into the export directory.
                 # convert_to_gguf writes .gguf files to cwd (repo root)
                 # because --outfile is a relative path like "model.Q4_K_M.gguf".
@@ -530,6 +563,11 @@ class ExportBackend:
 
                 self.current_model.save_pretrained(save_directory)
                 self.current_tokenizer.save_pretrained(save_directory)
+
+                # Fix broken tokenizer_class on disk (belt-and-suspenders)
+                from utils.transformers_version import patch_tokenizer_config
+                patch_tokenizer_config(save_directory, model_name=self._resolved_model_name)
+
                 logger.info(f"Adapter saved successfully to {save_directory}")
 
             # Push to hub if requested