Update trainer.py

trainer.py

@@ -187,6 +187,7 @@ if __name__ == "__main__":
     # read values from cfg as usual:
     conll_train_path = cfg["data"]["conll_train"]
     conll_dev_path   = cfg["data"].get("conll_dev")
+    conll_test_path   = cfg["data"].get("conll_test")  
     word_col_idx     = cfg["data"]["word_col_idx"]
     srl_first_col_idx= cfg["data"]["srl_first_col_idx"]
 
@@ -215,13 +216,19 @@ if __name__ == "__main__":
     tokenizer = AutoTokenizer.from_pretrained(replace_encoder_with or bert_name)
     print(f"Using tokenizer: {replace_encoder_with or bert_name}")
 
-    print(f"Loading multilingual CoNLL data: {conll_train_path}")
-    train_dataset, label2id, id2label = data_processing_for_loader_from_conll(
-        conll_path=conll_train_path,
-        tokenizer=tokenizer,
-        word_col_idx=word_col_idx,
-        srl_first_col_idx=srl_first_col_idx,
-    )
+    # print(f"Loading multilingual CoNLL data: {conll_train_path}")
+
+
+    train_bf_loader, dev_bf_loader, test_bf_loader, label2id, id2label = \
+        data_processing_for_loader_conll(
+            train_conll=conll_train_path,
+            dev_conll=conll_dev_path,
+            test_conll=conll_test_path,
+            tokenizer=tokenizer,
+            word_col_idx=word_col_idx,
+            srl_first_col_idx=srl_first_col_idx,
+            max_length=256,
+        )
 
     # pad_token_id = tokenizer.pad_token_id or tokenizer.eos_token_id
 
@@ -241,8 +248,9 @@ if __name__ == "__main__":
 
     collate = lambda b: srl_collate(b, pad_token_id=pad_token_id, pad_label_id=-100)
 
-    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, collate_fn=collate)
-    dev_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=False, collate_fn=collate)  # no dev split yet
+    train_loader = DataLoader(train_bf_loader, batch_size=batch_size, shuffle=True,  collate_fn=collate)
+    dev_loader   = DataLoader(dev_bf_loader,   batch_size=batch_size, shuffle=False, collate_fn=collate) if dev_bf_loader else None
+    test_loader  = DataLoader(test_bf_loader,  batch_size=batch_size, shuffle=False, collate_fn=collate) if test_bf_loader else None
 
     # ------------------------------
     # 🧠 Model initialization