#!/usr/bin/env python """ Simplified fine-tuning script for DeepSeek-R1-Distill-Qwen-14B-unsloth-bnb-4bit - Optimized for A100 GPU with pre-tokenized datasets - Research training only (no inference) - CLOUD BASED TRAINING - Hugging Face Spaces """ import os import logging import json import torch import argparse import shutil from pathlib import Path from datasets import load_dataset from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer, AutoConfig, BitsAndBytesConfig from transformers.data.data_collator import DataCollatorMixin from peft import LoraConfig, get_peft_model from dotenv import load_dotenv from huggingface_hub import HfApi, upload_folder # Basic environment setup for A100 os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True,max_split_size_mb:512" os.environ["NCCL_P2P_DISABLE"] = "1" # Can help with A100 multi-GPU setups # Force GPU mode in Space if we're using a pre-quantized model os.environ["FORCE_GPU"] = "1" # Disable tokenizers parallelism warning os.environ["TOKENIZERS_PARALLELISM"] = "false" # Create triton directory to avoid warning os.makedirs(os.path.expanduser("~/.triton/autotune"), exist_ok=True) # Default dataset with proper namespace DEFAULT_DATASET = "George-API/phi4-cognitive-dataset" # Set up logging logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') logger = logging.getLogger(__name__) # Determine if we're running in HF Space def is_running_in_space(): """Check if we're running in a Hugging Face Space""" return os.environ.get("SPACE_ID") is not None # Check if a model is pre-quantized (4-bit or 8-bit) def is_model_pre_quantized(model_name): """Check if model is already pre-quantized based on name""" pre_quantized_keywords = ["bnb-4bit", "4bit", "8bit", "quantized", "unsloth"] return any(keyword in model_name.lower() for keyword in pre_quantized_keywords) # Check if GPU is available def is_gpu_available(): """Simple check if CUDA is available according to PyTorch""" return torch.cuda.is_available() # Check if fully compatible CUDA is available for training def is_cuda_fully_available(model_name): """ Check if CUDA is fully available for training with bitsandbytes. More strict than torch.cuda.is_available() - requires full GPU compatibility. """ # If model is pre-quantized and we're in a Space with GPU selected, trust it if is_running_in_space() and is_model_pre_quantized(model_name) and is_gpu_available(): logger.info("Pre-quantized model detected with GPU in Hugging Face Space - using GPU mode") return True # For non-Space environments, or non-pre-quantized models, do detailed checks # If FORCE_GPU is set, trust that if os.environ.get("FORCE_GPU") == "1": logger.info("GPU mode forced by environment variable") return True # If running in Space and FORCE_GPU not explicitly set, be cautious if is_running_in_space() and os.environ.get("FORCE_GPU") != "1": # Check if CUDA is actually available if is_gpu_available(): logger.info("GPU detected in Hugging Face Space") return True else: logger.warning("No GPU detected in Hugging Face Space despite hardware selection") return False # If CUDA is not available according to PyTorch, we definitely can't use it if not is_gpu_available(): logger.warning("CUDA not available according to PyTorch") return False # Only test bitsandbytes if necessary (not for pre-quantized models) if not is_model_pre_quantized(model_name): try: import bitsandbytes as bnb logger.info("BitsAndBytes package is installed") # Try to create a dummy 4-bit computation to verify compatibility try: dummy = torch.zeros(1, device="cuda") a = bnb.nn.Linear4bit(1, 1) a.to(device="cuda") result = a(dummy) logger.info("BitsAndBytes with CUDA is working correctly") return True except Exception as e: logger.warning(f"BitsAndBytes CUDA compatibility test failed: {str(e)}") return False except ImportError: logger.warning("BitsAndBytes package not installed - cannot use 4-bit quantization") return False except Exception as e: logger.warning(f"Unexpected error checking BitsAndBytes: {str(e)}") return False # For pre-quantized models without bitsandbytes test return is_gpu_available() # Create a marker file to indicate training is active def create_training_marker(output_dir): os.makedirs(output_dir, exist_ok=True) with open("TRAINING_ACTIVE", "w") as f: f.write(f"Training active in {output_dir}") with open(os.path.join(output_dir, "RESEARCH_TRAINING_ONLY"), "w") as f: f.write("This model is for research training only. No interactive outputs.") # Remove the training marker file def remove_training_marker(): if os.path.exists("TRAINING_ACTIVE"): os.remove("TRAINING_ACTIVE") logger.info("Removed training active marker") # Function to upload model to Hugging Face Hub def upload_to_huggingface(output_dir, repo_name=None, private=False): """ Upload the trained model to Hugging Face Hub Args: output_dir: Directory containing the model files repo_name: Name of the repository on HF Hub (default: derived from output_dir) private: Whether the repository should be private (default: False) Returns: str: URL of the uploaded model on HF Hub """ logger.info(f"Uploading model from {output_dir} to Hugging Face Hub") # Get HF token from environment token = os.environ.get("HF_TOKEN") if not token: logger.error("HF_TOKEN environment variable not set. Please set it to upload to Hugging Face Hub.") logger.error("You can get a token from https://huggingface.co/settings/tokens") raise ValueError("HF_TOKEN not set") # Get or create repo name if not repo_name: # Use the output directory name as the repository name repo_name = os.path.basename(os.path.normpath(output_dir)) logger.info(f"Using repository name: {repo_name}") # Get HF username api = HfApi(token=token) user_info = api.whoami() username = user_info["name"] # Create full repository name full_repo_name = f"{username}/{repo_name}" logger.info(f"Creating repository: {full_repo_name}") # Create repository if it doesn't exist api.create_repo( repo_id=full_repo_name, exist_ok=True, private=private ) # Upload model files logger.info(f"Uploading files from {output_dir} to {full_repo_name}") api.upload_folder( folder_path=output_dir, repo_id=full_repo_name, commit_message="Upload model files" ) # Create model card model_card = f""" # {repo_name} This model was fine-tuned using the script at https://github.com/George-API/phi4-cognitive-dataset. ## Model details - Base model: DeepSeek-R1-Distill-Qwen-14B-unsloth-bnb-4bit - Dataset: {DEFAULT_DATASET} - Training: Research only """ with open(os.path.join(output_dir, "README.md"), "w") as f: f.write(model_card) # Upload the model card api.upload_file( path_or_fileobj=os.path.join(output_dir, "README.md"), path_in_repo="README.md", repo_id=full_repo_name, commit_message="Add model card" ) logger.info(f"Model successfully uploaded to https://huggingface.co/{full_repo_name}") return f"https://huggingface.co/{full_repo_name}" # Custom data collator for pre-tokenized data class PreTokenizedCollator(DataCollatorMixin): def __init__(self, pad_token_id=0, tokenizer=None): self.pad_token_id = pad_token_id self.tokenizer = tokenizer # Keep reference to tokenizer for fallback def __call__(self, features): # Extract features properly from the batch processed_features = [] for feature in features: # If input_ids is directly available, use it if 'input_ids' in feature and isinstance(feature['input_ids'], list): processed_features.append(feature) continue # If input_ids is not available, try to extract from conversations if 'input_ids' not in feature and 'conversations' in feature: conversations = feature['conversations'] if isinstance(conversations, list) and len(conversations) > 0: # Case 1: If conversations has 'input_ids' field (pre-tokenized) if isinstance(conversations[0], dict) and 'input_ids' in conversations[0]: feature['input_ids'] = conversations[0]['input_ids'] # Case 2: If conversations itself contains input_ids elif all(isinstance(x, int) for x in conversations): feature['input_ids'] = conversations # Case 3: If conversations has 'content' field elif isinstance(conversations[0], dict) and 'content' in conversations[0]: content = conversations[0]['content'] # If content is already tokens, use directly if isinstance(content, list) and all(isinstance(x, int) for x in content): feature['input_ids'] = content # If content is a string and we have tokenizer, tokenize as fallback elif isinstance(content, str) and self.tokenizer: logger.warning("Tokenizing string content as fallback") feature['input_ids'] = self.tokenizer.encode(content, add_special_tokens=False) # Ensure input_ids is present and is a list of integers if 'input_ids' in feature: if isinstance(feature['input_ids'], str) and self.tokenizer: feature['input_ids'] = self.tokenizer.encode(feature['input_ids'], add_special_tokens=False) elif not isinstance(feature['input_ids'], list): try: feature['input_ids'] = list(feature['input_ids']) except Exception as e: logger.error(f"Could not convert input_ids to list: {e}") continue processed_features.append(feature) if len(processed_features) == 0: raise ValueError("No valid examples found. Check dataset structure.") # Determine max length in this batch batch_max_len = max(len(x["input_ids"]) for x in processed_features) # Initialize batch tensors batch = { "input_ids": torch.ones((len(processed_features), batch_max_len), dtype=torch.long) * self.pad_token_id, "attention_mask": torch.zeros((len(processed_features), batch_max_len), dtype=torch.long), "labels": torch.ones((len(processed_features), batch_max_len), dtype=torch.long) * -100 # -100 is ignored in loss } # Fill batch tensors for i, feature in enumerate(processed_features): input_ids = feature["input_ids"] seq_len = len(input_ids) # Convert to tensor if it's a list if isinstance(input_ids, list): input_ids = torch.tensor(input_ids, dtype=torch.long) # Copy data to batch tensors batch["input_ids"][i, :seq_len] = input_ids batch["attention_mask"][i, :seq_len] = 1 # If there are labels, use them, otherwise use input_ids if "labels" in feature: labels = feature["labels"] if isinstance(labels, list): labels = torch.tensor(labels, dtype=torch.long) batch["labels"][i, :len(labels)] = labels else: batch["labels"][i, :seq_len] = input_ids return batch # Preprocess dataset to ensure all entries are pre-tokenized def preprocess_dataset(dataset, tokenizer): """Ensure dataset is fully pre-tokenized to avoid tokenization during training""" logger.info("Pre-processing dataset to ensure all entries are tokenized") def process_example(example): # If already has input_ids as list of integers, keep as is if 'input_ids' in example and isinstance(example['input_ids'], list) and all(isinstance(x, int) for x in example['input_ids']): return example # If has conversations with content field if 'conversations' in example: conversations = example['conversations'] if isinstance(conversations, list) and len(conversations) > 0: # If conversations has content field, tokenize it if isinstance(conversations[0], dict) and 'content' in conversations[0]: content = conversations[0]['content'] if isinstance(content, str): example['input_ids'] = tokenizer.encode(content, add_special_tokens=False) return example # For any other format, try to extract text and tokenize text = None if 'text' in example: text = example['text'] elif 'content' in example: text = example['content'] if text and isinstance(text, str): example['input_ids'] = tokenizer.encode(text, add_special_tokens=False) return example return dataset.map(process_example) # Load and prepare dataset with proper sorting def load_and_prepare_dataset(dataset_name, config, tokenizer=None): """Load and prepare the dataset for fine-tuning with proper sorting""" # Use the default dataset if the provided one matches the default name without namespace if dataset_name == "phi4-cognitive-dataset": dataset_name = DEFAULT_DATASET logger.info(f"Using full dataset path: {dataset_name}") logger.info(f"Loading dataset: {dataset_name}") try: # Load dataset try: dataset = load_dataset(dataset_name) except Exception as e: if "doesn't exist on the Hub or cannot be accessed" in str(e): logger.error(f"Dataset '{dataset_name}' not found. Make sure it exists and is accessible.") logger.error(f"If using a private dataset, check your HF_TOKEN is set in your environment.") logger.error(f"If missing namespace, try using the full path: 'George-API/phi4-cognitive-dataset'") raise # Extract the split we want to use (usually 'train') if 'train' in dataset: dataset = dataset['train'] # Get the dataset config dataset_config = config.get("dataset_config", {}) sort_field = dataset_config.get("sort_by_field", "prompt_number") # Preprocess dataset to ensure all entries are pre-tokenized if tokenizer is not None: dataset = preprocess_dataset(dataset, tokenizer) # Sort in ascending order by specified field logger.info(f"Sorting dataset by {sort_field} in ascending order") dataset = dataset.sort(sort_field) # Print dataset info logger.info(f"Dataset loaded with {len(dataset)} entries") logger.info(f"Dataset columns: {dataset.column_names}") # Print sample for debugging if len(dataset) > 0: logger.info(f"Sample entry structure: {list(dataset[0].keys())}") return dataset except Exception as e: logger.error(f"Error loading dataset: {str(e)}") raise # Load a simpler, smaller model for CPU mode def get_small_model_name(original_model_name): """Get a smaller model name for CPU mode""" # If using DeepSeek-R1-Distill-Qwen-14B, use a smaller model if "DeepSeek" in original_model_name and "14B" in original_model_name: logger.info("Using smaller model for CPU mode") return "distilgpt2" # Much smaller model # Otherwise just use the original model return original_model_name # Main training function def train(config_path, dataset_name, output_dir, upload_to_hub=False, hub_repo_name=None, private_repo=False): # Load environment variables load_dotenv() # Load config with open(config_path, 'r') as f: config = json.load(f) # Create training marker create_training_marker(output_dir) try: # Extract configs model_config = config.get("model_config", {}) training_config = config.get("training_config", {}) hardware_config = config.get("hardware_config", {}) lora_config = config.get("lora_config", {}) dataset_config = config.get("dataset_config", {}) # Log dataset info before loading logger.info(f"Will load dataset: {dataset_name}") if dataset_name != DEFAULT_DATASET and "phi4-cognitive-dataset" in dataset_name: logger.warning(f"Dataset name may need namespace prefix. Current: {dataset_name}") # Load model settings original_model_name = model_config.get("model_name_or_path") # Special handling for pre-quantized models like unsloth models is_pre_quantized = is_model_pre_quantized(original_model_name) if is_pre_quantized: logger.info(f"Detected pre-quantized model: {original_model_name}") # Determine if we can use CUDA with bitsandbytes can_use_4bit = is_cuda_fully_available(original_model_name) # For CPU mode, use a smaller model (unless pre-quantized) if not can_use_4bit and is_running_in_space() and not is_pre_quantized: model_name = get_small_model_name(original_model_name) logger.warning(f"Using smaller model {model_name} in CPU mode for Hugging Face Space") else: model_name = original_model_name logger.info(f"Using model: {model_name}") # Initialize tokenizer logger.info("Loading tokenizer") tokenizer = AutoTokenizer.from_pretrained( model_name, trust_remote_code=True ) tokenizer.pad_token = tokenizer.eos_token # Load and prepare dataset with proper sorting dataset = load_and_prepare_dataset(dataset_name, config, tokenizer) # Get quantization config quant_config = config.get("quantization_config", {}) # Determine if we should use 4-bit quantization # Pre-quantized models always use their built-in quantization if is_pre_quantized: use_4bit = True logger.info("Using pre-quantized model with built-in quantization") elif can_use_4bit and quant_config.get("load_in_4bit", True): use_4bit = True logger.info("Using 4-bit quantization with CUDA") else: use_4bit = False logger.warning("Using CPU mode without quantization") # Determine compute dtype based on hardware config compute_dtype = torch.bfloat16 if hardware_config.get("bf16", False) else torch.float16 logger.info(f"Using compute dtype: {compute_dtype}") # For pre-quantized models, always use device_map="auto" if is_pre_quantized and is_gpu_available(): logger.info("Loading pre-quantized model with GPU support") model = AutoModelForCausalLM.from_pretrained( model_name, device_map="auto", torch_dtype=compute_dtype, trust_remote_code=True, use_cache=model_config.get("use_cache", False) ) # Create model with proper configuration for non-pre-quantized models elif use_4bit and not is_pre_quantized: logger.info(f"Loading model with 4-bit quantization") # Create quantization config for GPU bnb_compute_dtype = torch.bfloat16 if quant_config.get("bnb_4bit_compute_dtype", "float16") == "bfloat16" else torch.float16 bnb_config = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_compute_dtype=bnb_compute_dtype, bnb_4bit_quant_type=quant_config.get("bnb_4bit_quant_type", "nf4"), bnb_4bit_use_double_quant=quant_config.get("bnb_4bit_use_double_quant", True) ) # Load 4-bit quantized model for GPU model = AutoModelForCausalLM.from_pretrained( model_name, quantization_config=bnb_config, device_map="auto", torch_dtype=compute_dtype, trust_remote_code=True, use_cache=model_config.get("use_cache", False), attn_implementation=hardware_config.get("attn_implementation", "flash_attention_2") ) else: # CPU fallback (or non-quantized GPU) mode logger.warning("Loading model in CPU fallback mode (no 4-bit quantization)") # Force CPU (safest option in HF Spaces) device_map = "cpu" dtype = torch.float32 logger.info("Forcing CPU mode for stability") # Load model without quantization model = AutoModelForCausalLM.from_pretrained( model_name, device_map=device_map, torch_dtype=dtype, trust_remote_code=True, use_cache=model_config.get("use_cache", False), low_cpu_mem_usage=True ) # Apply rope scaling if configured and available if "rope_scaling" in model_config and hasattr(model.config, "rope_scaling"): logger.info(f"Applying rope scaling: {model_config['rope_scaling']}") model.config.rope_scaling = model_config["rope_scaling"] # Create LoRA config logger.info("Creating LoRA configuration") # For pre-quantized models, we need proper target modules default_target_modules = ["q_proj", "k_proj", "v_proj", "o_proj"] # For pre-quantized models, especially Unsloth ones, we need to be careful with the target modules if is_pre_quantized: # For Unsloth models, use special configuration if "unsloth" in model_name.lower(): default_target_modules = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"] logger.info("Using Unsloth-specific LoRA target modules") lora_config_obj = LoraConfig( r=lora_config.get("r", 8), lora_alpha=lora_config.get("lora_alpha", 32), lora_dropout=lora_config.get("lora_dropout", 0.05), bias=lora_config.get("bias", "none"), task_type="CAUSAL_LM", # Explicitly set the task type target_modules=lora_config.get("target_modules", default_target_modules) ) # Apply LoRA to model logger.info("Applying LoRA to model") model = get_peft_model(model, lora_config_obj) logger.info("Successfully applied LoRA") # Ensure model parameters that need gradients are properly set if is_pre_quantized: logger.info("Verifying gradient settings for pre-quantized model") for name, param in model.named_parameters(): if 'lora' in name: # Only LoRA parameters should be trained if not param.requires_grad: logger.warning(f"LoRA parameter {name} doesn't have requires_grad=True, fixing...") param.requires_grad = True # Always use minimal batch size for HF Space CPU if is_running_in_space() and not can_use_4bit and not is_pre_quantized: per_device_train_batch_size = 1 logger.warning("Using minimal batch size for CPU training in Hugging Face Space") else: # Determine batch size based on available hardware if torch.cuda.is_available(): gpu_info = torch.cuda.get_device_properties(0) logger.info(f"GPU: {gpu_info.name}, VRAM: {gpu_info.total_memory / 1e9:.2f} GB") # Check if it's an A100 or high-memory GPU if "A100" in gpu_info.name or "A10G" in gpu_info.name or gpu_info.total_memory > 40e9: logger.info("Detected A100 GPU - optimizing for A100") per_device_train_batch_size = training_config.get("per_device_train_batch_size", 3) else: # Use a smaller batch size for other GPUs per_device_train_batch_size = 2 logger.info(f"Using conservative batch size for non-A100 GPU: {per_device_train_batch_size}") else: # Use minimal batch size for CPU per_device_train_batch_size = 1 logger.warning("No GPU detected - using minimal batch size for CPU training") # Use full training parameters for pre-quantized models or GPU mode if is_pre_quantized or can_use_4bit or not is_running_in_space(): num_train_epochs = training_config.get("num_train_epochs", 3) gradient_accumulation_steps = training_config.get("gradient_accumulation_steps", 2) fp16 = torch.cuda.is_available() and hardware_config.get("fp16", False) bf16 = torch.cuda.is_available() and hardware_config.get("bf16", True) # Disable gradient checkpointing for pre-quantized models as it can cause gradient issues gradient_checkpointing = torch.cuda.is_available() and hardware_config.get("gradient_checkpointing", True) and not is_pre_quantized dataloader_workers = training_config.get("dataloader_num_workers", 4) eval_strategy = training_config.get("eval_strategy", "no") load_best_model_at_end = False # Must be False when eval_strategy is "no" if is_pre_quantized: logger.info("Disabled gradient checkpointing for pre-quantized model to avoid gradient issues") logger.info("Using full training parameters for GPU mode") else: # For Space CPU training mode, use minimal parameters num_train_epochs = 1 gradient_accumulation_steps = 1 fp16 = False bf16 = False gradient_checkpointing = False dataloader_workers = 0 eval_strategy = "no" load_best_model_at_end = False logger.warning("Using minimal parameters for CPU training in Space") # Configure reporting backends reports = training_config.get("report_to", ["tensorboard"]) # Create training arguments logger.info("Creating training arguments") training_args = TrainingArguments( output_dir=output_dir, num_train_epochs=num_train_epochs, per_device_train_batch_size=per_device_train_batch_size, gradient_accumulation_steps=gradient_accumulation_steps, learning_rate=training_config.get("learning_rate", 2e-5), lr_scheduler_type=training_config.get("lr_scheduler_type", "cosine"), warmup_ratio=training_config.get("warmup_ratio", 0.03), weight_decay=training_config.get("weight_decay", 0.01), optim=training_config.get("optim", "adamw_torch"), fp16=fp16, bf16=bf16, max_grad_norm=training_config.get("max_grad_norm", 0.3), logging_steps=training_config.get("logging_steps", 10), save_steps=training_config.get("save_steps", 200), save_total_limit=training_config.get("save_total_limit", 3), eval_strategy=eval_strategy, load_best_model_at_end=load_best_model_at_end, report_to=reports, logging_first_step=training_config.get("logging_first_step", True), disable_tqdm=training_config.get("disable_tqdm", False), remove_unused_columns=False, gradient_checkpointing=gradient_checkpointing, dataloader_num_workers=dataloader_workers, group_by_length=training_config.get("group_by_length", True) ) # Create trainer with pre-tokenized collator logger.info("Creating trainer with pre-tokenized collator") trainer = Trainer( model=model, args=training_args, train_dataset=dataset, data_collator=PreTokenizedCollator( pad_token_id=tokenizer.pad_token_id, tokenizer=tokenizer ), # Add label_names to avoid warning compute_metrics=None, tokenizer=tokenizer # Provide tokenizer for proper padding ) # Start training logger.info("Starting training - RESEARCH PHASE ONLY") trainer.train() # Save the model logger.info(f"Saving model to {output_dir}") trainer.save_model(output_dir) # Save LoRA adapter separately lora_output_dir = os.path.join(output_dir, "lora_adapter") model.save_pretrained(lora_output_dir) logger.info(f"Saved LoRA adapter to {lora_output_dir}") # Save tokenizer tokenizer_output_dir = os.path.join(output_dir, "tokenizer") tokenizer.save_pretrained(tokenizer_output_dir) logger.info(f"Saved tokenizer to {tokenizer_output_dir}") # Save config for reference with open(os.path.join(output_dir, "training_config.json"), "w") as f: json.dump(config, f, indent=2) logger.info("Training complete - RESEARCH PHASE ONLY") # Upload to Hugging Face Hub if requested if upload_to_hub: hub_url = upload_to_huggingface( output_dir=output_dir, repo_name=hub_repo_name, private=private_repo ) logger.info(f"Model uploaded to Hugging Face Hub: {hub_url}") return output_dir finally: # Always remove the training marker when done remove_training_marker() if __name__ == "__main__": parser = argparse.ArgumentParser(description="Fine-tune DeepSeek model (Research Only)") parser.add_argument("--config", type=str, default="transformers_config.json", help="Path to the configuration file") parser.add_argument("--dataset", type=str, default=DEFAULT_DATASET, help="Dataset name or path") parser.add_argument("--output_dir", type=str, default="fine_tuned_model", help="Output directory for the fine-tuned model") parser.add_argument("--upload_to_hub", action="store_true", help="Upload the model to Hugging Face Hub after training") parser.add_argument("--hub_repo_name", type=str, default=None, help="Repository name for the model on Hugging Face Hub") parser.add_argument("--private_repo", action="store_true", help="Make the Hugging Face Hub repository private") parser.add_argument("--force_cpu", action="store_true", help="Force CPU mode even if CUDA is available") args = parser.parse_args() # Force CPU mode if requested if args.force_cpu: os.environ["FORCE_GPU"] = "0" logger.info("Forcing CPU mode as requested") try: output_path = train( args.config, args.dataset, args.output_dir, upload_to_hub=args.upload_to_hub, hub_repo_name=args.hub_repo_name, private_repo=args.private_repo ) print(f"Research training completed. Model saved to: {output_path}") if args.upload_to_hub: print("Model was also uploaded to Hugging Face Hub.") except Exception as e: logging.error(f"Training failed: {str(e)}") remove_training_marker() # Clean up marker if training fails raise