specifying batch_size did not work - using training argument specification per HF docs and ChatGPT

epi_pipeline.py

@@ -426,7 +426,7 @@ from torch import nn
 import numpy as np
 from unidecode import unidecode
 import re
-from transformers import BertConfig, AutoModelForTokenClassification, BertTokenizer, Trainer
+from transformers import BertConfig, AutoModelForTokenClassification, BertTokenizer, Trainer, TrainingArguments
 from unidecode import unidecode
 from collections import OrderedDict
 import json
@@ -655,7 +655,13 @@ class NER_Pipeline:
         #model = AutoModelForTokenClassification.from_pretrained(name_or_path_to_model_folder)
         self.config = BertConfig.from_pretrained(name_or_path_to_model_folder)
         self.labels = {re.sub(".-","",label) for label in self.config.label2id.keys() if label != "O"}
-        self.trainer = Trainer(model=AutoModelForTokenClassification.from_pretrained(name_or_path_to_model_folder))
+        # Create the TrainingArguments object and set the batch size
+        training_args = TrainingArguments(
+            per_device_eval_batch_size=16,
+            per_gpu_eval_batch_size=16,
+            # other training arguments...
+            )
+        self.trainer = Trainer(model=AutoModelForTokenClassification.from_pretrained(name_or_path_to_model_folder),args=training_args)
     
     def __str__(self):
         return "Instantiation: pipe = NER_Pipeline(name_or_path_to_model_folder)"+"\n Calling: output_dict = pipe(text)"
@@ -669,7 +675,7 @@ class NER_Pipeline:
         output_dict = {label:[] for label in self.labels}
         
         dataset = NerDataset(text, self.bert_tokenizer, self.config)
-        predictions, label_ids, _ = self.trainer.predict(dataset, batch_size=16)
+        predictions, label_ids, _ = self.trainer.predict(dataset)
         preds_list, _ = self.align_predictions(predictions, label_ids)
         #dataset.ner_inputs.labels = preds_list
         for ner_input, sent_pred_list in zip(dataset.ner_inputs, preds_list):