Update train_model.py

train_model.py

@@ -1,5 +1,5 @@
 # train_model.py (Training Script)
-from huggingface_hub import login
+
 import argparse
 from transformers import (
     GPT2Config,
@@ -15,7 +15,7 @@ from transformers import (
 from datasets import load_dataset, Dataset
 import torch
 import os
-from huggingface_hub import HfApi, HfFolder
+from huggingface_hub import login, HfApi, HfFolder
 import logging
 
 def setup_logging(log_file_path):
@@ -49,7 +49,7 @@ def parse_arguments():
     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 (e.g., 'username/dataset')")
+    parser.add_argument("--dataset_name", type=str, required=True, help="Name of the Hugging Face dataset (e.g., 'wikitext/wikitext-2-raw-v1')")
     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")
@@ -65,14 +65,21 @@ def load_and_prepare_dataset(task, dataset_name, tokenizer, sequence_length):
     logging.info(f"Loading dataset '{dataset_name}' for task '{task}'...")
     try:
         if task == "generation":
-            train_dataset = load_dataset(dataset_name,split='train',use_auth_token=True)
-            dataset = train_dataset['train'].shuffle(seed=42).select(range(500))
+            # Check if dataset_name includes config
+            if '/' in dataset_name:
+                dataset, config = dataset_name.split('/', 1)
+                dataset = load_dataset(dataset, config, split='train[:1%]', use_auth_token=True)
+            else:
+                dataset = load_dataset(dataset_name, split='train[:1%]', use_auth_token=True)
             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":
-            train_dataset = load_dataset(dataset_name,split='train',use_auth_token=True)
-            dataset = train_dataset['train'].shuffle(seed=42).select(range(500))
+            if '/' in dataset_name:
+                dataset, config = dataset_name.split('/', 1)
+                dataset = load_dataset(dataset, config, split='train[:1%]', use_auth_token=True)
+            else:
+                dataset = load_dataset(dataset_name, split='train[:1%]', use_auth_token=True)
             logging.info("Dataset loaded successfully for classification task.")
             # Assuming the dataset has 'text' and 'label' columns
             def tokenize_function(examples):
@@ -80,7 +87,8 @@ def load_and_prepare_dataset(task, dataset_name, tokenizer, sequence_length):
         else:
             raise ValueError("Unsupported task type")
         
-        tokenized_datasets = dataset.map(tokenize_function, batched=True)
+        # Shuffle and select a subset
+        tokenized_datasets = dataset.shuffle(seed=42).select(range(500)).map(tokenize_function, batched=True)
         logging.info("Dataset tokenization complete.")
         return tokenized_datasets
     except Exception as e:
@@ -139,18 +147,22 @@ def main():
 
     # Initialize Hugging Face API
     api = HfApi()
-    hf_token = os.getenv('HF_API_TOKEN')
+    
+    # Retrieve the Hugging Face API token from environment variables
+    hf_token = os.getenv("HF_API_TOKEN")
     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
+        logging.error("HF_API_TOKEN environment variable not set.")
+        raise ValueError("HF_API_TOKEN environment variable not set.")
+    
+    # Perform login using the API token
     try:
         login(token=hf_token)
         logging.info("Successfully logged in to Hugging Face Hub.")
     except Exception as e:
         logging.error(f"Failed to log in to Hugging Face Hub: {str(e)}")
         raise e
+
+    # Initialize tokenizer
     try:
         logging.info("Initializing tokenizer...")
         if args.task == "generation":