fixed inference errors

Author: Ubuntu

src/bart_reddit_lora/evaluation.py

@@ -1,87 +1,65 @@
 """
-Metrics computation module for sequence-to-sequence models.
-
-This module provides a factory function to create a `compute_metrics` callable
-for Hugging Face's `Trainer`. The returned function computes ROUGE-L, BLEU, and
-BERTScore (F1) on decoded model predictions versus labels.
+Module for building a ROUGE-L metric computation function
+for Hugging Face Seq2SeqTrainer.
 """
 
 import numpy as np
 import evaluate
-from transformers import EvalPrediction
-from typing import Callable, Dict, Any, Union
-from transformers.tokenization_utils_base import PreTrainedTokenizerBase
+from transformers import Seq2SeqTrainer, Seq2SeqTrainingArguments, EarlyStoppingCallback, PreTrainedTokenizerBase, EvalPrediction
+from typing import Callable, Dict
 
 
-def build_compute_metrics(
-    tok: PreTrainedTokenizerBase, 
-    num_process_workers: int = 2
-) -> Callable[[EvalPrediction], Dict[str, float]]:
+def build_compute_metrics(tok: PreTrainedTokenizerBase) -> Callable[[EvalPrediction], Dict[str, float]]:
     """
-    Create a metrics computation function for use with Hugging Face `Trainer`.
+    Create a compute_metrics function for Seq2SeqTrainer that returns the ROUGE-L score.
 
     Args:
-        tokenizer: A Hugging Face tokenizer for decoding predictions/labels.
-        num_process_workers: Number of worker processes for metric computation.
+        tok (PreTrainedTokenizerBase): Tokenizer for decoding predictions and labels.
 
     Returns:
-        A callable that takes an `EvalPrediction` and returns a dict with:
-          - "rougeL": ROUGE-L score (%)
-          - "bleu": BLEU score (%)
-          - "bertscore_f1": average BERTScore F1
+        Callable[[EvalPrediction], Dict[str, float]]: Function computing "rougeL" percentage.
     """
-    rouge = evaluate.load("rouge")  # longest-substring overlap
-    bleu = evaluate.load("bleu")  # n-gram precision
-    bertscore = evaluate.load("bertscore")  # semantic similarity
+    rouge = evaluate.load("rouge", keep_in_memory=True)   # keep_in_memory avoids disk I/O
 
-    def _compute_metrics(eval_pred: EvalPrediction) -> Dict[str, float]:
+    # 2️⃣  Metric fn: decode → strip → compute → return only rougeL
+    def compute_metrics(eval_pred):
         """
-        Compute ROUGE-L, BLEU, and BERTScore given model predictions and labels.
+        Decode predictions and references, compute ROUGE-L, and return as percentage.
 
         Args:
-            eval_pred: An `EvalPrediction` with `predictions` and `label_ids`.
+            eval_pred (EvalPrediction): Object with .predictions and .label_ids.
 
         Returns:
-            A dict mapping metric names to rounded scores.
+            Dict[str, float]: Dictionary with key "rougeL" and its percentage score.
         """
-        preds, labels = eval_pred.predictions, eval_pred.label_ids
-
-        # handle tuple output (some models return (generated_ids, ...))
-        if isinstance(preds, tuple):
+        preds, labels = eval_pred
+        if isinstance(preds, tuple):          # HF sometimes returns (logits, ...)
             preds = preds[0]
 
-        # decode
-        decoded_preds = tok.batch_decode(preds, skip_special_tokens=True)
+        # Replace label pad tokens (-100) so they can be decoded
         labels = np.where(labels != -100, labels, tok.pad_token_id)
-        decoded_labels = tok.batch_decode(labels, skip_special_tokens=True)
 
-        # metrics
-        rouge_l = rouge.compute(
+        decoded_preds  = tok.batch_decode(preds, skip_special_tokens=True,
+                                        clean_up_tokenization_spaces=True)
+        decoded_labels = tok.batch_decode(labels, skip_special_tokens=True,
+                                        clean_up_tokenization_spaces=True)
+
+        # Strip white-space/newlines that can hurt ROUGE scores
+        decoded_preds  = [s.strip() for s in decoded_preds]
+        decoded_labels = [s.strip() for s in decoded_labels]
+
+        score_dict = rouge.compute(
             predictions=decoded_preds,
             references=decoded_labels,
-            use_stemmer=True,
-            num_process_workers=num_process_workers,
-        )["rougeL"]
-        bleu_score = bleu.compute(
-            predictions=decoded_preds,
-            references=[[ref] for ref in decoded_labels],  # BLEU expects list-of-lists
-            smooth=True,
-            num_process_workers=num_process_workers,
-        )["bleu"]
-        bert_f1 = np.mean(
-            bertscore.compute(
-                predictions=decoded_preds,
-                references=decoded_labels,
-                lang="en",
-                num_process_workers=num_process_workers,
-            )["f1"]
+            use_stemmer=True,        # standard setting for ROUGE-* in HF evaluate
         )
 
-        # round for nice logging
-        return {
-            "rougeL": round(rouge_l * 100, 4),
-            "bleu": round(bleu_score * 100, 4),
-            "bertscore_f1": round(bert_f1, 4),
-        }
+        # HF’s rouge.compute() returns fractional scores; convert to %
+        rougeL = round(score_dict["rougeL"] * 100, 4)
+
+        return {"rougeL": rougeL}
+
+    return compute_metrics
+
+
 
-    return _compute_metrics

src/bart_reddit_lora/inference.py

@@ -35,7 +35,6 @@ class InferenceArgs:
         mode: Either 'test' to evaluate on the test dataset or 'predict' to generate outputs for raw texts.
         batch_size: Batch size used for evaluation or prediction.
         texts: List of input strings when running in 'predict' mode.
-        num_process_workers: Number of processes for parallel metric computation.
         use_sdpa_attention: Whether to enable SDPA attention for memory efficiency.
     """
     mode: str = field(
@@ -51,10 +50,6 @@ class InferenceArgs:
         default_factory=list,
         metadata={"help": "One or more input texts for `predict` mode."},
     )
-    num_process_workers: int = field(
-        default=2,
-        metadata={"help": "Number of workers to parallelize n-gram counting."},
-    )
     use_sdpa_attention: bool = field(
         default=True, metadata={"help": "Enable Sdpa for mem-efficient kernel."}
     )
@@ -114,19 +109,33 @@ def main() -> None:
             args=Seq2SeqTrainingArguments(
                 output_dir="outputs/inference",
                 per_device_eval_batch_size=inf_args.batch_size,
-                predict_with_generate=True,
-                generation_max_length=384,
+                predict_with_generate=False,
+                # generation_max_length=640,
                 report_to=[],
             ),
             eval_dataset=ds_tok["test"],
             data_collator=data_collator,
             tokenizer=tok,
-            compute_metrics=build_compute_metrics(tok, inf_args.num_process_workers),
+            # compute_metrics=build_compute_metrics(tok),
         )
-
-        pred_output = trainer.predict(ds_tok["test"])
-        metrics = pred_output.metrics
+        metrics = trainer.evaluate(ds_tok["test"])
         logger.info(f"Test metrics: {metrics}")
+        
+        test_loader = trainer.get_eval_dataloader()
+        model = trainer.model
+        device = trainer.args.device
+
+        losses = []
+        with torch.no_grad():
+            for batch in test_loader:
+                # move inputs → device
+                batch = {k: v.to(device) for k, v in batch.items()}
+                # forward pass: passing in labels returns `loss`
+                outputs = model(**batch)
+                losses.append(outputs.loss.item())
+
+        mean_loss = sum(losses) / len(losses)
+        logger.info(f"Test loss: {mean_loss:.4f}")
 
     elif inf_args.mode == "predict":
         if not inf_args.texts:

src/bart_reddit_lora/train.py

@@ -63,7 +63,7 @@ class CustomTrainingArgs(Seq2SeqTrainingArguments):
     num_train_epochs: int = 12
     per_device_train_batch_size: int = 32
     per_device_eval_batch_size: int = 64
-    learning_rate: float = 3e-5
+    learning_rate: float = 6e-5
     lr_scheduler_type: str = "cosine"
     warmup_ratio: float = 0.1
     max_grad_norm: float = 0.5