Update app.py

app.py

@@ -38,6 +38,19 @@ def summarize():
     data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, model=checkpoint)
    
     rouge = evaluate.load("rouge")
+
+    def compute_metrics(eval_pred):
+        predictions, labels = eval_pred
+        decoded_preds = tokenizer.batch_decode(predictions, skip_special_tokens=True)
+        labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
+        decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
+
+        result = rouge.compute(predictions=decoded_preds, references=decoded_labels, use_stemmer=True)
+
+        prediction_lens = [np.count_nonzero(pred != tokenizer.pad_token_id) for pred in predictions]
+        result["gen_len"] = np.mean(prediction_lens)
+
+        return {k: round(v, 4) for k, v in result.items()}
     
     return data_collator
 #    return type(tokenized_billsum)