Update train.py

train.py

@@ -1,4 +1,4 @@
-from transformers import MBartForSequenceClassification, BertTokenizer, TrainingArguments, Trainer
+from transformers import BertModel, BertTokenizer, TrainingArguments, Trainer
 from datasets import Dataset
 
 
@@ -28,7 +28,7 @@ def init():
     # Load the model and tokenizer
     model_name = "language-ml-lab/AzerBert"  # Replace with your model name if different
     tokenizer = BertTokenizer.from_pretrained(model_name)
-    model = MBartForSequenceClassification.from_pretrained(model_name)
+    model = BertModel.from_pretrained(model_name)
     chunk_size = 512
 
     # Load the plain text (replace with your actual loading logic)