Update train.py

train.py

@@ -36,7 +36,7 @@ def init():
       constitution_text = f.read()
     
     # Prepare the dataset
-    train_dataset = prepare_text_dataset(constitution_text)
+    train_dataset = prepare_text_dataset(constitution_text, chunk_size)
     
     # Define training arguments
     training_args = TrainingArguments(
@@ -56,7 +56,7 @@ def init():
     )
     
     # Start training
-    trainer.train(chunk_size)
+    trainer.train()
     
     # Save the fine-tuned model
     model.save_pretrained("./fine-tuned_model")