Upload folder using huggingface_hub

run_transformers_training.py

@@ -8,6 +8,16 @@ import argparse
 import logging
 from datetime import datetime
 
+# Import Unsloth first, before other ML imports
+try:
+    from unsloth import FastLanguageModel
+    from unsloth.chat_templates import get_chat_template
+    unsloth_available = True
+except ImportError:
+    unsloth_available = False
+    logger = logging.getLogger(__name__)
+    logger.warning("Unsloth not available. Please install with: pip install unsloth")
+
 import torch
 from datasets import load_dataset
 from transformers import (
@@ -44,15 +54,6 @@ except ImportError:
     peft_available = False
     logger.warning("PEFT not available. Parameter-efficient fine-tuning will not be used.")
 
-# Import Unsloth
-try:
-    from unsloth import FastLanguageModel
-    from unsloth.chat_templates import get_chat_template
-    unsloth_available = True
-except ImportError:
-    unsloth_available = False
-    logger.warning("Unsloth not available. Please install with: pip install unsloth")
-
 def load_env_variables():
     """Load environment variables from system, .env file, or Hugging Face Space variables."""
     # Check if we're running in a Hugging Face Space
@@ -131,6 +132,13 @@ def load_model_and_tokenizer(config):
             logger.error("Unsloth is required for training with pre-quantized model")
             logger.error("Please ensure unsloth is in requirements.txt")
             raise ImportError("Unsloth is required for this training setup")
+        
+        # Get model name correctly from nested config structure
+        model_name = config.get("model", {}).get("name") or config.get("model_name_or_path") or config.get("model_name")
+        logger.info(f"Loading model: {model_name}")
+        
+        if not model_name:
+            raise ValueError("Model name not found in configuration. Please check your transformers_config.json file.")
             
         logger.info("Using Unsloth optimizations with pre-quantized model")
         # Check for flash attention without importing it directly
@@ -143,32 +151,37 @@ def load_model_and_tokenizer(config):
             logger.warning("Flash attention not available, falling back to standard attention")
             
         model, tokenizer = FastLanguageModel.from_pretrained(
-            model_name=config.get("model_name"),
-            max_seq_length=config.get("max_seq_length", 2048),
+            model_name=model_name,
+            max_seq_length=config.get("max_seq_length", 2048) or config.get("tokenizer", {}).get("max_seq_length", 2048),
             dtype=None,  # Let Unsloth choose optimal dtype
             device_map="auto",
             # Don't explicitly use flash attention config here, let Unsloth handle it
         )
         
         # Apply Unsloth's training optimizations with config parameters
+        unsloth_config = config.get("unsloth", {})
         model = FastLanguageModel.get_peft_model(
             model,
-            r=config.get("unsloth_r", 32),
-            target_modules=config.get("unsloth_target_modules", 
+            r=unsloth_config.get("r", 32),
+            target_modules=unsloth_config.get("target_modules", 
                 ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]),
-            lora_alpha=config.get("unsloth_alpha", 16),
-            lora_dropout=config.get("unsloth_dropout", 0.05),
+            lora_alpha=unsloth_config.get("alpha", 16),
+            lora_dropout=unsloth_config.get("dropout", 0.05),
             bias="none",
-            use_gradient_checkpointing=config.get("gradient_checkpointing", True),
+            use_gradient_checkpointing=config.get("gradient_checkpointing", True) or config.get("training", {}).get("gradient_checkpointing", True),
             random_state=config.get("seed", 42),
         )
         logger.info("Unsloth optimizations applied successfully")
 
         # Set up tokenizer settings
-        if config.get("chat_template"):
-            chat_template = get_chat_template("phi")
-            tokenizer.chat_template = chat_template
-            logger.info("Set phi chat template")
+        chat_template = config.get("chat_template") or config.get("tokenizer", {}).get("chat_template")
+        if chat_template:
+            try:
+                template = get_chat_template("phi")
+                tokenizer.chat_template = template
+                logger.info("Set phi chat template")
+            except Exception as e:
+                logger.warning(f"Failed to set chat template: {str(e)}")
         
         # Ensure proper token settings
         if tokenizer.pad_token_id is None:
@@ -418,24 +431,58 @@ def main():
     # Load all configurations
     try:
         configs = load_configs(args.config_dir)
+        
+        # Extract specific configs
+        if not configs:
+            logger.error("Failed to load configurations")
+            return 1
+            
+        # Verify configurations exist
+        if "transformers" not in configs:
+            logger.error("transformers_config.json not found or invalid")
+            return 1
+            
+        if "hardware" not in configs:
+            logger.warning("hardware_config.json not found. Using default hardware configuration.")
+            
+        if "dataset" not in configs:
+            logger.error("dataset_config.json not found or invalid")
+            return 1
+            
+        # Validate model configuration
+        model_config = configs["transformers"]
+        if not model_config.get("model", {}).get("name") and not model_config.get("model_name_or_path") and not model_config.get("model_name"):
+            logger.error("Model name not specified in configuration")
+            logger.error("Please ensure 'name' is specified under 'model' in transformers_config.json")
+            return 1
+            
+        logger.info(f"Model name: {model_config.get('model', {}).get('name') or model_config.get('model_name_or_path') or model_config.get('model_name')}")
         logger.info("All configurations loaded successfully")
         
         # Extract specific configs
         model_config = configs["transformers"]
-        hardware_config = configs["hardware"]
+        hardware_config = configs.get("hardware", {})
         dataset_config = configs["dataset"]
         
-        # Apply hardware-specific settings
-        per_device_batch_size = hardware_config["training_optimizations"]["per_device_batch_size"]
-        gradient_accumulation = hardware_config["training_optimizations"]["gradient_accumulation_steps"]
-        
-        # Update model config with hardware settings
-        model_config["training"].update({
-            "per_device_train_batch_size": per_device_batch_size,
-            "gradient_accumulation_steps": gradient_accumulation,
-            "gradient_checkpointing": hardware_config["training_optimizations"]["memory_optimizations"]["use_gradient_checkpointing"]
-        })
-        
+        # Apply hardware-specific settings if available
+        if hardware_config:
+            training_opts = hardware_config.get("training_optimizations", {})
+            per_device_batch_size = training_opts.get("per_device_batch_size")
+            gradient_accumulation = training_opts.get("gradient_accumulation_steps")
+            
+            if per_device_batch_size and model_config.get("training"):
+                model_config["training"]["per_device_train_batch_size"] = per_device_batch_size
+                logger.info(f"Applied hardware-specific batch size: {per_device_batch_size}")
+                
+            if gradient_accumulation and model_config.get("training"):
+                model_config["training"]["gradient_accumulation_steps"] = gradient_accumulation
+                logger.info(f"Applied hardware-specific gradient accumulation: {gradient_accumulation}")
+                
+            # Apply memory optimizations
+            memory_opts = training_opts.get("memory_optimizations", {})
+            if memory_opts.get("use_gradient_checkpointing") is not None and model_config.get("training"):
+                model_config["training"]["gradient_checkpointing"] = memory_opts["use_gradient_checkpointing"]
+                
     except Exception as e:
         logger.error(f"Error loading configurations: {e}")
         return 1
@@ -445,75 +492,16 @@ def main():
     set_seed(seed)
     logger.info(f"Set random seed to {seed}")
     
-    # Check if we're running in a Hugging Face Space
-    if os.environ.get("SPACE_ID") and not os.environ.get("HF_USERNAME"):
-        # Extract username from SPACE_ID
-        username = os.environ.get("SPACE_ID").split("/")[0]
-        logger.info(f"Extracted username from SPACE_ID: {username}")
-        
-        # Set hub_model_id if not already set and push_to_hub is enabled
-        if model_config.get("push_to_hub", False) and not model_config.get("hub_model_id"):
-            model_name = model_config.get("model_name", "").split("/")[-1]
-            model_config["hub_model_id"] = f"{username}/finetuned-{model_name}"
-            logger.info(f"Set hub_model_id to {model_config['hub_model_id']}")
-    
-    # Load model and tokenizer
-    logger.info(f"Loading model: {model_config.get('model_name')}")
-    
     try:
         model, tokenizer = load_model_and_tokenizer(model_config)
         logger.info("Model and tokenizer loaded successfully")
         
-        # Prepare model for k-bit training if using PEFT
-        if model_config.get("use_peft", False) and peft_available:
-            logger.info("Preparing model for parameter-efficient fine-tuning")
-            try:
-                model = prepare_model_for_kbit_training(model)
-                
-                # Get target modules
-                target_modules = model_config.get("target_modules", ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"])
-                
-                # Create LoRA config
-                lora_config = LoraConfig(
-                    r=model_config.get("lora_r", 16),
-                    lora_alpha=model_config.get("lora_alpha", 32),
-                    lora_dropout=model_config.get("lora_dropout", 0.05),
-                    bias="none",
-                    task_type="CAUSAL_LM",
-                    target_modules=target_modules
-                )
-                
-                # Apply LoRA to model
-                model = get_peft_model(model, lora_config)
-                logger.info(f"Applied LoRA with r={model_config.get('lora_r', 16)}, alpha={model_config.get('lora_alpha', 32)}")
-            except Exception as e:
-                logger.error(f"Error setting up PEFT: {e}")
-                return 1
-        
-        # Load dataset
-        logger.info(f"Loading dataset: {dataset_config.get('dataset_name')}")
+        # Load dataset with proper mapping
         try:
-            dataset = load_dataset(dataset_config.get("dataset_name"))
-            logger.info(f"Dataset loaded successfully with {len(dataset['train'])} training examples")
-            
-            # Sort dataset by ID to ensure chunks from the same paper are processed together
-            logger.info("Sorting dataset by ID to maintain paper chunk order")
-            def sort_by_id(example):
-                # Extract ID as integer if possible, otherwise keep as string
-                try:
-                    return int(example['id'])
-                except (ValueError, TypeError):
-                    return example['id']
-            
-            # Apply sorting to the dataset
-            dataset['train'] = dataset['train'].sort('id')
-            logger.info("Dataset sorted by ID")
-            
-            # Log the first few IDs to verify sorting
-            sample_ids = [example['id'] for example in dataset['train'].select(range(min(5, len(dataset['train']))))]
-            logger.info(f"First few IDs after sorting: {sample_ids}")
+            dataset = load_dataset_with_mapping(dataset_config)
+            logger.info("Dataset loaded and prepared successfully")
         except Exception as e:
-            logger.error(f"Error loading or sorting dataset: {e}")
+            logger.error(f"Error loading dataset: {e}")
             return 1
         
         # Create data collator