Update train_model.py

train_model.py

@@ -14,18 +14,20 @@ from huggingface_hub import HfApi, HfFolder
 import logging
 
 def main():
-    # Configure Logging
-    logging.basicConfig(
-        filename='training.log',
-        filemode='a',
-        format='%(asctime)s - %(levelname)s - %(message)s',
-        level=logging.INFO
-    )
+    # ... 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 ...
+
 
     parser = argparse.ArgumentParser()
     parser.add_argument("--task", type=str, required=True, help="Task type: generation or classification")
     parser.add_argument("--model_name", type=str, required=True, help="Name of the model")
-    parser.add_argument("--dataset", type=str, required=True, help="Path to the dataset")
+    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)
@@ -53,26 +55,17 @@ def main():
     
     # Load and prepare dataset
     if args.task == "generation":
-        dataset = load_dataset('text', data_files={'train': args.dataset})
+        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":
-        # For classification, assume the dataset is a simple text file with "text\tlabel" per line
-        with open(args.dataset, "r", encoding="utf-8") as f:
-            lines = f.readlines()
-        texts = []
-        labels = []
-        for line in lines:
-            parts = line.strip().split("\t")
-            if len(parts) == 2:
-                texts.append(parts[0])
-                labels.append(int(parts[1]))
-        dataset = Dataset.from_dict({"text": texts, "label": labels})
+        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")
-    
+
     tokenized_datasets = dataset.map(tokenize_function, batched=True)
     
     if args.task == "generation":