From 6f6c590e5798b6aa0c37df5be1655bfd19b3eeca Mon Sep 17 00:00:00 2001
From: theblackcat102 <theblackcat102@github.com>
Date: Sat, 14 Jan 2023 06:47:21 +0000
Subject: [PATCH] [fix] Disable task specific evaluation

---
 .../custom_datasets/summarization.py          |  1 +
 model/supervised_finetuning/utils.py          | 57 ++++++++++---------
 2 files changed, 30 insertions(+), 28 deletions(-)

diff --git a/model/supervised_finetuning/custom_datasets/summarization.py b/model/supervised_finetuning/custom_datasets/summarization.py
index 41fa6dc0..69e4b51d 100644
--- a/model/supervised_finetuning/custom_datasets/summarization.py
+++ b/model/supervised_finetuning/custom_datasets/summarization.py
@@ -8,6 +8,7 @@ SUMMARIZATION_SPECIAL_TOKENS = {"Text": "", "Summary": ["TL;DR:", "Summarize thi
 SUMMARY_SPECIAL_PROMPT = {
     "multi_news": ["Summarize in bullet points", "Generate summary in list of points"],
     "xsum": ["Give me summary in one sentence", "Short TLDR", "Give me a concise summary"],
+    "samsum": ["TLDR;", "Summarize this dialogue", "Summarize dialogue"],
 }
 
 summarization_config_mapping = {
diff --git a/model/supervised_finetuning/utils.py b/model/supervised_finetuning/utils.py
index 85fb86db..7b6e03b6 100644
--- a/model/supervised_finetuning/utils.py
+++ b/model/supervised_finetuning/utils.py
@@ -1,12 +1,13 @@
-from functools import partial
+# from functools import partial
 from pathlib import Path
 
 import evaluate
-import nltk
-import numpy as np
+
+# import nltk
+# import numpy as np
 import transformers
 import yaml
-from custom_datasets import QA_DATASETS, SUMMARIZATION_DATASETS, get_one_dataset
+from custom_datasets import get_one_dataset
 from custom_datasets.dialogue_collator import DialogueDataCollator
 from custom_datasets.qa_datasets import QA_SPECIAL_TOKENS
 from losses import CrossEntropyLoss, PolyLoss
@@ -52,25 +53,25 @@ def preprocess_qa(eval_pred):
     return (eval_pred.predictions, eval_pred.label_ids)
 
 
-def postprocess_summarization(preds, labels):
-    preds = [pred.strip() for pred in preds]
-    labels = [label.strip() for label in labels]
+# def postprocess_summarization(preds, labels):
+#     preds = [pred.strip() for pred in preds]
+#     labels = [label.strip() for label in labels]
 
-    preds = ["\n".join(nltk.sent_tokenize(pred)) for pred in preds]
-    labels = ["\n".join(nltk.sent_tokenize(label)) for label in labels]
+#     preds = ["\n".join(nltk.sent_tokenize(pred)) for pred in preds]
+#     labels = ["\n".join(nltk.sent_tokenize(label)) for label in labels]
 
-    return preds, labels
+#     return preds, labels
 
 
-def preprocess_summarization(eval_pred, tokenizer, ignore_pad_token_for_loss=True):
-    preds, labels = eval_pred
-    decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)
-    if ignore_pad_token_for_loss:
-        labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
-    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
+# def preprocess_summarization(eval_pred, tokenizer, ignore_pad_token_for_loss=True):
+#     preds, labels = eval_pred
+#     decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)
+#     if ignore_pad_token_for_loss:
+#         labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
+#     decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
 
-    decoded_preds, decoded_labels = postprocess_summarization(decoded_preds, decoded_labels)
-    return decoded_preds, decoded_labels
+#     decoded_preds, decoded_labels = postprocess_summarization(decoded_preds, decoded_labels)
+#     return decoded_preds, decoded_labels
 
 
 def get_metrics(conf, tokenizer):
@@ -78,16 +79,16 @@ def get_metrics(conf, tokenizer):
     # metrics in the future for more thorough evaluation
     metrics, preprocess_fns = [evaluate.load("accuracy")], [default_preprocess]
 
-    if any(dataset in QA_DATASETS for dataset in conf.datasets):
-        raise ValueError("TODO")
-        metrics.append(evaluate.load("squad_v2"))
-        preprocess_fns.append(preprocess_qa)
-    if any(dataset in SUMMARIZATION_DATASETS for dataset in conf.datasets):
-        raise ValueError("TODO")
-        metrics.append(evaluate.load("rouge"))
-        preprocess_fns.append(
-            partial(preprocess_summarization, tokenizer, ignore_pad_token_for_loss=conf.ignore_pad_token_for_loss)
-        )
+    # if any(dataset in QA_DATASETS for dataset in conf.datasets):
+    #     raise ValueError("TODO")
+    #     metrics.append(evaluate.load("squad_v2"))
+    #     preprocess_fns.append(preprocess_qa)
+    # if any(dataset in SUMMARIZATION_DATASETS for dataset in conf.datasets):
+    #     raise ValueError("TODO")
+    #     metrics.append(evaluate.load("rouge"))
+    #     preprocess_fns.append(
+    #         partial(preprocess_summarization, tokenizer, ignore_pad_token_for_loss=conf.ignore_pad_token_for_loss)
+    #     )
 
     return metrics, preprocess_fns