Remove _attn_implementation_autoset if present in config.json (#60)

fixie-ai · Feb 6, 2025 · bcd42c8 · bcd42c8
1 parent 7e8ea5e
commit bcd42c8
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 0 deletions.
diff --git a/ultravox/model/ultravox_config.py b/ultravox/model/ultravox_config.py
@@ -169,7 +169,12 @@ def to_diff_dict(self) -> Dict[str, Any]:
         # remove text_config and audio_config if text_model_id and audio_model_id are present
         if self.text_model_id is not None:
             diff_dict.pop("text_config", None)
+        elif "text_config" in diff_dict:
+            diff_dict["text_config"].pop("_attn_implementation_autoset", None)
+
         if self.audio_model_id is not None:
             diff_dict.pop("audio_config", None)
+        elif "audio_config" in diff_dict:
+            diff_dict["audio_config"].pop("_attn_implementation_autoset", None)
 
         return diff_dict
diff --git a/ultravox/ultravoxls/ultravoxls_config.py b/ultravox/ultravoxls/ultravoxls_config.py
@@ -100,5 +100,7 @@ def to_diff_dict(self) -> Dict[str, Any]:
         # remove text_config if text_model_id
         if self.text_model_id is not None:
             diff_dict.pop("text_config", None)
+        elif "text_config" in diff_dict:
+            diff_dict["text_config"].pop("_attn_implementation_autoset", None)
 
         return diff_dict