Use sacrebleu directly for training metrics

src/train_forward.py

@@ -2,7 +2,7 @@
 Forward synthesis model training script.
 Trains T5 model to predict products from reactants.
 """
-import evaluate
+import sacrebleu
 import numpy as np
 from transformers import (
     AutoTokenizer,
@@ -71,23 +71,21 @@ def main():
     collator = DataCollatorForSeq2Seq(tokenizer, model=model, padding=True)
     
     # Metrics
-    metric = evaluate.load("sacrebleu")
-    
     def compute_metrics(eval_pred):
         preds, labels = eval_pred
         preds = np.where(preds != -100, preds, tokenizer.pad_token_id)
         labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
-        
+
         decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)
         decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
-        
+
         decoded_preds = [p.strip() for p in decoded_preds]
         decoded_labels = [l.strip() for l in decoded_labels]
-        
-        bleu = metric.compute(predictions=decoded_preds, references=[[l] for l in decoded_labels])
+
+        bleu = sacrebleu.corpus_bleu(decoded_preds, [decoded_labels])
         exact = np.mean([p == l for p, l in zip(decoded_preds, decoded_labels)])
-        
-        return {"bleu": bleu["score"], "exact_match": exact}
+
+        return {"bleu": bleu.score, "exact_match": exact}
     
     # Trainer
     print("\nInitializing trainer...")

src/train_retro.py

@@ -2,7 +2,7 @@
 Retrosynthesis model training script.
 Trains T5 model to predict reactants from products.
 """
-import evaluate
+import sacrebleu
 import numpy as np
 from transformers import (
     AutoTokenizer,
@@ -71,23 +71,21 @@ def main():
     collator = DataCollatorForSeq2Seq(tokenizer, model=model, padding=True)
     
     # Metrics
-    metric = evaluate.load("sacrebleu")
-    
     def compute_metrics(eval_pred):
         preds, labels = eval_pred
         preds = np.where(preds != -100, preds, tokenizer.pad_token_id)
         labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
-        
+
         decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)
         decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
-        
+
         decoded_preds = [p.strip() for p in decoded_preds]
         decoded_labels = [l.strip() for l in decoded_labels]
-        
-        bleu = metric.compute(predictions=decoded_preds, references=[[l] for l in decoded_labels])
+
+        bleu = sacrebleu.corpus_bleu(decoded_preds, [decoded_labels])
         exact = np.mean([p == l for p, l in zip(decoded_preds, decoded_labels)])
-        
-        return {"bleu": bleu["score"], "exact_match": exact}
+
+        return {"bleu": bleu.score, "exact_match": exact}
     
     # Trainer
     print("\nInitializing trainer...")