Update train.py

train.py

@@ -1,16 +1,26 @@
 from datasets import load_dataset
-from transformers import Trainer, TrainingArguments, Tacotron2ForConditionalGeneration
+from transformers import FastSpeechForConditionalGeneration, Trainer, TrainingArguments
 
-# تحميل البيانات من Hugging Face Datasets
+# تحميل البيانات للهجة النجدية
 dataset = load_dataset("m6011/sada2022")
 najdi_data = dataset.filter(lambda example: example['SpeakerDialect'] == 'Najdi')
 
-# إعداد النموذج والمعالج
-model = Tacotron2ForConditionalGeneration.from_pretrained("facebook/tacotron2")
+# إعداد النموذج
+model = FastSpeechForConditionalGeneration.from_pretrained("facebook/fastspeech2-en-ljspeech")
 
-# إعداد التدريب
-training_args = TrainingArguments(output_dir="./results", per_device_train_batch_size=16, num_train_epochs=3)
-trainer = Trainer(model=model, args=training_args, train_dataset=najdi_data)
+# إعداد المدرب
+training_args = TrainingArguments(
+    output_dir="./results",
+    per_device_train_batch_size=4,
+    num_train_epochs=5,
+)
+
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=najdi_data['train'],
+    eval_dataset=najdi_data['test']
+)
 
 # بدء التدريب
 trainer.train()