BusinessDev commited on
Commit
b606f8a
1 Parent(s): 2eec429

Update train.py

Browse files
Files changed (1) hide show
  1. train.py +2 -2
train.py CHANGED
@@ -1,4 +1,4 @@
1
- from transformers import MBartForSequenceClassification, GPT2Tokenizer, TrainingArguments, Trainer
2
  from datasets import Dataset
3
 
4
 
@@ -27,7 +27,7 @@ def prepare_text_dataset(data, chunk_size):
27
  def init():
28
  # Load the model and tokenizer
29
  model_name = "language-ml-lab/AzerBert" # Replace with your model name if different
30
- tokenizer = GPT2Tokenizer.from_pretrained(model_name)
31
  model = MBartForSequenceClassification.from_pretrained(model_name)
32
  chunk_size = 512
33
 
 
1
+ from transformers import MBartForSequenceClassification, BertTokenizer, TrainingArguments, Trainer
2
  from datasets import Dataset
3
 
4
 
 
27
  def init():
28
  # Load the model and tokenizer
29
  model_name = "language-ml-lab/AzerBert" # Replace with your model name if different
30
+ tokenizer = BertTokenizer.from_pretrained(model_name)
31
  model = MBartForSequenceClassification.from_pretrained(model_name)
32
  chunk_size = 512
33