Update train_model.py

train_model.py

@@ -2,10 +2,15 @@
 
 import argparse
 from transformers import (
-    GPT2Config, GPT2LMHeadModel, 
-    BertConfig, BertForSequenceClassification,
-    Trainer, TrainingArguments, AutoTokenizer,
-    DataCollatorForLanguageModeling, DataCollatorWithPadding
+    GPT2Config,
+    GPT2LMHeadModel,
+    BertConfig,
+    BertForSequenceClassification,
+    Trainer,
+    TrainingArguments,
+    AutoTokenizer,
+    DataCollatorForLanguageModeling,
+    DataCollatorWithPadding,
 )
 from datasets import load_dataset, Dataset
 import torch
@@ -13,110 +18,196 @@ import os
 from huggingface_hub import HfApi, HfFolder
 import logging
 
-def main():
-    # ... existing code ...
-    if args.task == "generation":
-        dataset = load_dataset(args.dataset_name, split='train')  # Load dataset by name
-    elif args.task == "classification":
-        dataset = load_dataset(args.dataset_name, split='train')  # Adjust if necessary
-    else:
-        raise ValueError("Unsupported task type")
-    # ... existing code ...
+def setup_logging(log_file_path):
+    """
+    Sets up logging to both console and a file.
+    """
+    logger = logging.getLogger()
+    logger.setLevel(logging.INFO)
+
+    # Create handlers
+    c_handler = logging.StreamHandler()
+    f_handler = logging.FileHandler(log_file_path)
+    c_handler.setLevel(logging.INFO)
+    f_handler.setLevel(logging.INFO)
 
+    # Create formatters and add to handlers
+    c_format = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
+    f_format = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
+    c_handler.setFormatter(c_format)
+    f_handler.setFormatter(f_format)
 
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--task", type=str, required=True, help="Task type: generation or classification")
+    # Add handlers to the logger
+    logger.addHandler(c_handler)
+    logger.addHandler(f_handler)
+
+def parse_arguments():
+    """
+    Parses command-line arguments.
+    """
+    parser = argparse.ArgumentParser(description="Train a custom LLM.")
+    parser.add_argument("--task", type=str, required=True, choices=["generation", "classification"],
+                        help="Task type: 'generation' or 'classification'")
     parser.add_argument("--model_name", type=str, required=True, help="Name of the model")
-    parser.add_argument("--dataset_name", type=str, required=True, help="Name of the Hugging Face dataset")
-    parser.add_argument("--num_layers", type=int, default=12)
-    parser.add_argument("--attention_heads", type=int, default=1)
-    parser.add_argument("--hidden_size", type=int, default=64)
-    parser.add_argument("--vocab_size", type=int, default=30000)
-    parser.add_argument("--sequence_length", type=int, default=512)
+    parser.add_argument("--dataset_name", type=str, required=True, help="Name of the Hugging Face dataset (e.g., 'username/dataset')")
+    parser.add_argument("--num_layers", type=int, default=12, help="Number of hidden layers")
+    parser.add_argument("--attention_heads", type=int, default=1, help="Number of attention heads")
+    parser.add_argument("--hidden_size", type=int, default=64, help="Hidden size of the model")
+    parser.add_argument("--vocab_size", type=int, default=30000, help="Vocabulary size")
+    parser.add_argument("--sequence_length", type=int, default=512, help="Maximum sequence length")
     args = parser.parse_args()
-    
-    logging.info(f"Starting training for model: {args.model_name}, Task: {args.task}")
+    return args
+
+def load_and_prepare_dataset(task, dataset_name, tokenizer, sequence_length):
+    """
+    Loads and tokenizes the dataset based on the task.
+    """
+    logging.info(f"Loading dataset '{dataset_name}' for task '{task}'...")
+    try:
+        if task == "generation":
+            dataset = load_dataset(dataset_name, split='train')
+            logging.info("Dataset loaded successfully for generation task.")
+            def tokenize_function(examples):
+                return tokenizer(examples['text'], truncation=True, max_length=sequence_length)
+        elif task == "classification":
+            dataset = load_dataset(dataset_name, split='train')
+            logging.info("Dataset loaded successfully for classification task.")
+            # Assuming the dataset has 'text' and 'label' columns
+            def tokenize_function(examples):
+                return tokenizer(examples['text'], truncation=True, max_length=sequence_length)
+        else:
+            raise ValueError("Unsupported task type")
+        
+        tokenized_datasets = dataset.map(tokenize_function, batched=True)
+        logging.info("Dataset tokenization complete.")
+        return tokenized_datasets
+    except Exception as e:
+        logging.error(f"Error loading or tokenizing dataset: {str(e)}")
+        raise e
+
+def initialize_model(task, model_name, vocab_size, sequence_length, hidden_size, num_layers, attention_heads):
+    """
+    Initializes the model configuration and model based on the task.
+    """
+    logging.info(f"Initializing model for task '{task}'...")
+    try:
+        if task == "generation":
+            config = GPT2Config(
+                vocab_size=vocab_size,
+                n_positions=sequence_length,
+                n_ctx=sequence_length,
+                n_embd=hidden_size,
+                num_hidden_layers=num_layers,
+                num_attention_heads=attention_heads,
+                intermediate_size=4 * hidden_size,
+                hidden_act='gelu',
+                use_cache=True
+            )
+            model = GPT2LMHeadModel(config)
+            logging.info("GPT2LMHeadModel initialized successfully.")
+        elif task == "classification":
+            config = BertConfig(
+                vocab_size=vocab_size,
+                max_position_embeddings=sequence_length,
+                hidden_size=hidden_size,
+                num_hidden_layers=num_layers,
+                num_attention_heads=attention_heads,
+                intermediate_size=4 * hidden_size,
+                hidden_act='gelu',
+                num_labels=2  # Adjust based on your classification task
+            )
+            model = BertForSequenceClassification(config)
+            logging.info("BertForSequenceClassification initialized successfully.")
+        else:
+            raise ValueError("Unsupported task type")
+        
+        return model
+    except Exception as e:
+        logging.error(f"Error initializing model: {str(e)}")
+        raise e
+
+def main():
+    # Parse arguments
+    args = parse_arguments()
+
+    # Setup logging
+    log_file = "training.log"
+    setup_logging(log_file)
+    logging.info("Training script started.")
 
-    # Define output directory
-    output_dir = f"./models/{args.model_name}"
-    os.makedirs(output_dir, exist_ok=True)
-    
     # Initialize Hugging Face API
     api = HfApi()
     hf_token = HfFolder.get_token()
-    
-    # Initialize tokenizer (adjust based on task)
-    if args.task == "generation":
-        tokenizer = AutoTokenizer.from_pretrained("gpt2")
-    elif args.task == "classification":
-        tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
-    else:
-        raise ValueError("Unsupported task type")
-    
+    if not hf_token:
+        logging.error("HF_API_TOKEN is not set. Please set it as an environment variable.")
+        raise ValueError("HF_API_TOKEN is not set.")
+
+    # Initialize tokenizer
+    try:
+        logging.info("Initializing tokenizer...")
+        if args.task == "generation":
+            tokenizer = AutoTokenizer.from_pretrained("gpt2")
+        elif args.task == "classification":
+            tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
+        else:
+            raise ValueError("Unsupported task type")
+        logging.info("Tokenizer initialized successfully.")
+    except Exception as e:
+        logging.error(f"Error initializing tokenizer: {str(e)}")
+        raise e
+
     # Load and prepare dataset
-    if args.task == "generation":
-        dataset = load_dataset(args.dataset_name, split='train')
-        def tokenize_function(examples):
-            return tokenizer(examples['text'], truncation=True, max_length=args.sequence_length)
-    elif args.task == "classification":
-        dataset = load_dataset(args.dataset_name, split='train')
-        # Assuming the dataset has 'text' and 'label' columns
-        def tokenize_function(examples):
-            return tokenizer(examples['text'], truncation=True, max_length=args.sequence_length)
-    else:
-        raise ValueError("Unsupported task type")
+    try:
+        tokenized_datasets = load_and_prepare_dataset(
+            task=args.task,
+            dataset_name=args.dataset_name,
+            tokenizer=tokenizer,
+            sequence_length=args.sequence_length
+        )
+    except Exception as e:
+        logging.error("Failed to load and prepare dataset.")
+        raise e
 
-    tokenized_datasets = dataset.map(tokenize_function, batched=True)
-    
+    # Initialize model
+    try:
+        model = initialize_model(
+            task=args.task,
+            model_name=args.model_name,
+            vocab_size=args.vocab_size,
+            sequence_length=args.sequence_length,
+            hidden_size=args.hidden_size,
+            num_layers=args.num_layers,
+            attention_heads=args.attention_heads
+        )
+    except Exception as e:
+        logging.error("Failed to initialize model.")
+        raise e
+
+    # Define data collator
     if args.task == "generation":
         data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
     elif args.task == "classification":
         data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
-    
-    # Initialize model based on task
-    if args.task == "generation":
-        config = GPT2Config(
-            vocab_size=args.vocab_size,
-            n_positions=args.sequence_length,
-            n_ctx=args.sequence_length,
-            n_embd=args.hidden_size,
-            num_hidden_layers=args.num_layers,
-            num_attention_heads=args.attention_heads,
-            intermediate_size=4 * args.hidden_size,
-            hidden_act='gelu',
-            use_cache=True
-        )
-        model = GPT2LMHeadModel(config)
-    elif args.task == "classification":
-        config = BertConfig(
-            vocab_size=args.vocab_size,
-            max_position_embeddings=args.sequence_length,
-            hidden_size=args.hidden_size,
-            num_hidden_layers=args.num_layers,
-            num_attention_heads=args.attention_heads,
-            intermediate_size=4 * args.hidden_size,
-            hidden_act='gelu',
-            num_labels=2  # Adjust based on your classification task
-        )
-        model = BertForSequenceClassification(config)
     else:
-        raise ValueError("Unsupported task type")
-    
+        logging.error("Unsupported task type for data collator.")
+        raise ValueError("Unsupported task type for data collator.")
+
     # Define training arguments
     if args.task == "generation":
         training_args = TrainingArguments(
-            output_dir=output_dir,
+            output_dir=f"./models/{args.model_name}",
             num_train_epochs=3,
             per_device_train_batch_size=8,
             save_steps=5000,
             save_total_limit=2,
             logging_steps=500,
             learning_rate=5e-4,
-            remove_unused_columns=False
+            remove_unused_columns=False,
+            push_to_hub=False  # We'll handle pushing manually
         )
     elif args.task == "classification":
         training_args = TrainingArguments(
-            output_dir=output_dir,
+            output_dir=f"./models/{args.model_name}",
             num_train_epochs=3,
             per_device_train_batch_size=16,
             evaluation_strategy="epoch",
@@ -124,37 +215,61 @@ def main():
             save_total_limit=2,
             logging_steps=500,
             learning_rate=5e-5,
-            remove_unused_columns=False
+            remove_unused_columns=False,
+            push_to_hub=False  # We'll handle pushing manually
         )
-    
+    else:
+        logging.error("Unsupported task type for training arguments.")
+        raise ValueError("Unsupported task type for training arguments.")
+
     # Initialize Trainer
     trainer = Trainer(
         model=model,
         args=training_args,
-        train_dataset=tokenized_datasets['train'],
+        train_dataset=tokenized_datasets,
         data_collator=data_collator,
     )
-    
+
     # Start training
-    trainer.train()
-    
-    # Save the final model
-    trainer.save_model(output_dir)
-    tokenizer.save_pretrained(output_dir)
-    
-    # Push to Hugging Face Hub
-    model_repo = f"your-username/{args.model_name}"  # Replace 'your-username' with your actual username
+    logging.info("Starting training...")
     try:
+        trainer.train()
+        logging.info("Training completed successfully.")
+    except Exception as e:
+        logging.error(f"Error during training: {str(e)}")
+        raise e
+
+    # Save the final model and tokenizer
+    try:
+        trainer.save_model(training_args.output_dir)
+        tokenizer.save_pretrained(training_args.output_dir)
+        logging.info(f"Model and tokenizer saved to '{training_args.output_dir}'.")
+    except Exception as e:
+        logging.error(f"Error saving model or tokenizer: {str(e)}")
+        raise e
+
+    # Push the model to Hugging Face Hub
+    model_repo = f"{api.whoami(token=hf_token)['name']}/{args.model_name}"
+    try:
+        logging.info(f"Pushing model to Hugging Face Hub at '{model_repo}'...")
         api.create_repo(repo_id=model_repo, private=False, token=hf_token)
+        logging.info(f"Repository '{model_repo}' created successfully.")
     except Exception as e:
-        logging.warning(f"Repository might already exist: {e}")
-    model.push_to_hub(model_repo, use_auth_token=hf_token)
-    tokenizer.push_to_hub(model_repo, use_auth_token=hf_token)
-    
-    logging.info(f"Model '{args.model_name}' trained and pushed to Hugging Face Hub at '{model_repo}'.")
+        logging.warning(f"Repository might already exist: {str(e)}")
     
+    try:
+        model.push_to_hub(model_repo, use_auth_token=hf_token)
+        tokenizer.push_to_hub(model_repo, use_auth_token=hf_token)
+        logging.info(f"Model and tokenizer pushed to Hugging Face Hub at '{model_repo}'.")
+    except Exception as e:
+        logging.error(f"Error pushing model to Hugging Face Hub: {str(e)}")
+        raise e
+
+    logging.info("Training script finished successfully.")
+
 if __name__ == "__main__":
     main()
 
 
 
+