Upload train_aviation.py with huggingface_hub

train_aviation.py

@@ -3,7 +3,7 @@
 #     "torch", 
 #     "trl>=0.12.0",
 #     "peft>=0.7.0",
-#     "transformers>=4.46.0",
+#     "transformers", # Let UV pick latest
 #     "huggingface_hub>=0.26.0",
 #     "accelerate>=0.24.0",
 #     "trackio",
@@ -42,39 +42,7 @@ from peft import LoraConfig, prepare_model_for_kbit_training, get_peft_model
 from trl import SFTTrainer, SFTConfig
 from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, AutoConfig
 
-# Register 'ministral3' config to handle nested text_config (Removed, as patching directly)
-# This whole section is being replaced by direct config patching below.
-# print("🔧 Registering ministral3 config (Monkey Patch Strategy)...")
-# try:
-#     from transformers import MinistralConfig, AutoConfig
-#     class Ministral3CompatConfig(MinistralConfig):
-#         model_type = "ministral3" # Ensure this matches the `text_config["model_type"]`
-#         def __init__(self, **kwargs):
-#             super().__init__(**kwargs)
-#             if not hasattr(self, 'sliding_window') or self.sliding_window is None:
-#                 self.sliding_window = 4096
-#             if not hasattr(self, 'layer_types'):
-#                 self.layer_types = ["sliding_attention"] * getattr(self, "num_hidden_layers", 40)
-#     AutoConfig.register("ministral3", Ministral3CompatConfig)
-#     print("   Registered ministral3 -> Ministral3CompatConfig (patched)")
-# except Exception as e:
-#     print(f"   ❌ Failed to patch/register ministral3 config: {e}")
-
-# Register Mistral3Config to a model class (Removed, not needed with direct patching)
-# print("🔧 Registering Mistral3 model class...")
-# try:
-#     from transformers.models.mistral3.configuration_mistral3 import Mistral3Config
-#     try:
-#         from transformers.models.mistral3.modeling_mistral3 import Mistral3ForConditionalGeneration
-#         AutoModelForCausalLM.register(Mistral3Config, Mistral3ForConditionalGeneration)
-#         print("   Registered Mistral3Config -> Mistral3ForConditionalGeneration")
-#     except ImportError:
-#         print("   Mistral3ForConditionalGeneration not found, trying MistralForCausalLM")
-#         from transformers import MistralForCausalLM
-#         AutoModelForCausalLM.register(Mistral3Config, MistralForCausalLM)
-#         print("   Registered Mistral3Config -> MistralForCausalLM")
-# except ImportError as e:
-#     print(f"   ❌ Failed to find Mistral3Config or register model: {e}")
+# All custom config registration logic removed, relying on latest transformers
 
 
 # Load dataset
@@ -120,27 +88,27 @@ bnb_config = BitsAndBytesConfig(
     bnb_4bit_use_double_quant=True,
 )
 
-# Load config first to patch
+# Load config first (AutoConfig should handle it now with latest transformers)
 print(f"🤖 Loading config for {model_id}...")
 config = AutoConfig.from_pretrained(model_id)
 
-# Patch text_config to include sliding_window and layer_types
-print("🔧 Patching config.text_config...")
-if hasattr(config, 'text_config'):
-    if not hasattr(config.text_config, 'sliding_window') or config.text_config.sliding_window is None:
-        config.text_config.sliding_window = 4096
-        print("   Set config.text_config.sliding_window = 4096")
-    if not hasattr(config.text_config, 'layer_types'):
-        config.text_config.layer_types = ["sliding_attention"] * getattr(config.text_config, "num_hidden_layers", 40)
-        print("   Set config.text_config.layer_types")
-else:
-    print("   No text_config found, skipping patching.")
-
-# Load Model with the patched config
-print(f"🤖 Loading model {model_id} with patched config...")
+# Patch text_config to include sliding_window and layer_types (Now unnecessary, should be handled by latest transformers)
+# print("🔧 Patching config.text_config...")
+# if hasattr(config, 'text_config'):
+#     if not hasattr(config.text_config, 'sliding_window') or config.text_config.sliding_window is None:
+#         config.text_config.sliding_window = 4096
+#         print("   Set config.text_config.sliding_window = 4096")
+#     if not hasattr(config.text_config, 'layer_types'):
+#         config.text_config.layer_types = ["sliding_attention"] * getattr(config.text_config, "num_hidden_layers", 40)
+#         print("   Set config.text_config.layer_types")
+# else:
+#     print("   No text_config found, skipping patching.")
+
+# Load Model with the config
+print(f"🤖 Loading model {model_id} with config...")
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
-    config=config, # Pass the patched config
+    config=config, # Pass the config
     quantization_config=bnb_config,
     device_map="auto",
     torch_dtype=torch.bfloat16,
@@ -203,4 +171,4 @@ print("🚀 Starting training...")
 trainer.train()
 
 print("💾 Pushing to Hub...")
-trainer.push_to_hub()
+trainer.push_to_hub()