Update train.py

train.py

@@ -1,8 +1,8 @@
 from datasets import load_dataset
 from transformers import AutoTokenizer, DataCollatorForLanguageModeling, Trainer, TrainingArguments, AutoModelForMaskedLM
 
-# Load dataset from local CSV
-dataset = load_dataset("text", data_files="chunks.csv")
+# Load dataset from Hugging Face Hub
+dataset = load_dataset("drzeeIslam/nelson-gpt-chunks")
 
 # Load tokenizer and model
 model_checkpoint = "distilbert-base-uncased"
@@ -36,5 +36,5 @@ trainer = Trainer(
     data_collator=data_collator
 )
 
-# Train the model
+# Start training
 trainer.train()