Update train.py

train.py

@@ -4,7 +4,7 @@ from datasets import Dataset
 
 
 # Prepare the dataset (simplified)
-def prepare_text_dataset(data):
+def prepare_text_dataset(data, chunk_size):
     # Split the text into smaller chunks (consider logical divisions of the Constitution)
     chunks = [data[i:i+chunk_size] for i in range(0, len(data), chunk_size)]
     # Convert chunks to dictionaries with a single feature "text"
@@ -56,7 +56,7 @@ def init():
     )
     
     # Start training
-    trainer.train()
+    trainer.train(chunk_size)
     
     # Save the fine-tuned model
     model.save_pretrained("./fine-tuned_model")