Upload train_aviation.py with huggingface_hub

train_aviation.py

@@ -18,13 +18,15 @@ import os
 from huggingface_hub import list_repo_files
 
 # DEBUG: Check token and repo access
+# (commented out for cleaner logs now that it works)
 # print("🔍 DIAGNOSTICS:")
 # token = os.environ.get("HF_TOKEN")
 # print(f"   HF_TOKEN env var present: {bool(token)}")
 # if token:
 #     print(f"   HF_TOKEN prefix: {token[:4]}...")
 
-# model_id = "mistralai/Ministral-3-14B-Reasoning-2512"
+model_id = "mistralai/Ministral-3-14B-Reasoning-2512" # Defined at top level
+
 # try:
 #     print(f"   Attempting to list files for {model_id}...")
 #     files = list_repo_files(model_id, token=token)
@@ -34,61 +36,45 @@ from huggingface_hub import list_repo_files
 #     print(f"   ❌ Failed to list repo files: {e}")
 # print("="*40)
 
-model_id = "mistralai/Ministral-3-14B-Reasoning-2512"
 
 from datasets import load_dataset
 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
-print("🔧 Registering ministral3 config (Monkey Patch Strategy)...")
-try:
-    from transformers import MinistralConfig, AutoConfig
-    
-    # We need to ensure MinistralConfig has sliding_window and layer_types if it's used
-    # as the inner text_config for Mistral3.
-    # Create a temporary compatible class.
-    
-    class Ministral3CompatConfig(MinistralConfig):
-        model_type = "ministral3" # Ensure this matches the `text_config["model_type"]`
-        def __init__(self, **kwargs):
-            super().__init__(**kwargs)
-            # Ensure sliding_window is set, if null in config.json or missing
-            if not hasattr(self, 'sliding_window') or self.sliding_window is None:
-                self.sliding_window = 4096 # Default value for Mistral/Ministral models
-
-            # Ensure layer_types is set, as it's expected by modeling_ministral.py
-            if not hasattr(self, 'layer_types'):
-                # Assumes all layers are sliding attention if the model uses it
-                # Use getattr for num_hidden_layers as it might not be set yet if config is partial
-                self.layer_types = ["sliding_attention"] * getattr(self, "num_hidden_layers", 40) # Default to 40 if not found
-    
-    # Register the compatible class for the "ministral3" key
-    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
-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}")
-
-# Model ID
-# model_id defined above
+# 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}")
 
 
 # Load dataset
@@ -134,10 +120,27 @@ bnb_config = BitsAndBytesConfig(
     bnb_4bit_use_double_quant=True,
 )
 
-# Load Model
-print(f"🤖 Loading model {model_id}...")
+# Load config first to patch
+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...")
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
+    config=config, # Pass the patched config
     quantization_config=bnb_config,
     device_map="auto",
     torch_dtype=torch.bfloat16,
@@ -200,4 +203,4 @@ print("🚀 Starting training...")
 trainer.train()
 
 print("💾 Pushing to Hub...")
-trainer.push_to_hub()
+trainer.push_to_hub()