[feature] added translation, rallio instruct tuning dataset, prosocial for safety, new summary dataset

model/supervised_finetuning/custom_datasets/summarization.py

@@ -1,3 +1,6 @@
+"""
+    Summarize different spectrum of documents
+"""
 import random
 
 from datasets import load_dataset
@@ -12,13 +15,21 @@ SUMMARY_SPECIAL_PROMPT = {
 }
 
 summarization_config_mapping = {
-    "cnn_dailymail": ("3.0.0",),
-    "samsum": (),
-    "xsum": (),
-    "multi_news": (),
-    "scitldr": ("AIC",),
-    "billsum": (),
-    "reddit": (),
+    "cnn_dailymail": (
+        "cnn_dailymail",
+        "3.0.0",
+    ),
+    "samsum": ("samsum",),
+    "xsum": ("xsum",),
+    "multi_news": ("multi_news",),
+    "scitldr": (
+        "scitldr",
+        "AIC",
+    ),
+    "billsum": ("billsum",),
+    "reddit": ("reddit",),
+    "tldr_news": ("JulesBelveze/tldr_news",),  # need to fix : JulesBelveze/tldr_news
+    "debate_sum": ("Hellisotherpeople/DebateSum",),  # Hellisotherpeople/DebateSum
 }
 
 summarization_name_mapping = {
@@ -29,6 +40,8 @@ summarization_name_mapping = {
     "scitldr": ("source", "target"),
     "billsum": ("text", "summary"),
     "reddit": ("content", "summary"),
+    "tldr_news": ("content", "headline"),
+    "debate_sum": ("Full-Document", "Extract"),
 }
 
 
@@ -43,7 +56,7 @@ def index_summary_merge(text, summary):
 class SummarizationDataset(Dataset):
     def __init__(self, dataset, cache_dir, split):
         self.name = dataset
-        self.dataset = load_dataset(dataset, *summarization_config_mapping[dataset], cache_dir=cache_dir, split=split)
+        self.dataset = load_dataset(*summarization_config_mapping[dataset], cache_dir=cache_dir, split=split)
         self.text_column, self.summary_column = summarization_name_mapping[dataset]
         self.preprocess_fn = index_summary_merge if dataset == "scitldr" else index_summary_default