mirror of
https://github.com/wassname/Open-Assistant.git
synced 2026-07-05 17:30:48 +08:00
theblackcat102
58 lines
1.7 KiB
Python
58 lines
1.7 KiB
Python
import random
|
|
from datasets import load_dataset
|
|
from torch.utils.data import Dataset
|
|
|
|
SUMMARIZATION_SPECIAL_TOKENS = {"Text": "", "Summary": ["TL;DR:", "Summarize this", "Give me the summary"]}
|
|
|
|
|
|
|
|
summarization_config_mapping = {
|
|
"cnn_dailymail": ("3.0.0",),
|
|
"samsum": (),
|
|
"xsum": (),
|
|
"multi_news": (),
|
|
"scitldr": ("AIC",),
|
|
"billsum": (),
|
|
"reddit": (),
|
|
}
|
|
|
|
summarization_name_mapping = {
|
|
"cnn_dailymail": ("article", "highlights"),
|
|
"samsum": ("dialogue", "summary"),
|
|
"xsum": ("document", "summary"),
|
|
"multi_news": ("document", "summary"),
|
|
"scitldr": ("source", "target"),
|
|
"billsum": ("text", "summary"),
|
|
"reddit": ("content", "summary"),
|
|
}
|
|
|
|
|
|
def index_summary_default(text, summary):
|
|
return text.replace('\n\n', '\n'), summary
|
|
|
|
|
|
def index_summary_merge(text, summary):
|
|
return " ".join(text), " ".join(summary)
|
|
|
|
|
|
class SummarizationDataset(Dataset):
|
|
def __init__(self, dataset, cache_dir, split):
|
|
self.dataset = load_dataset(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 == "scitdlr" else index_summary_default
|
|
|
|
def __len__(self):
|
|
return len(self.dataset)
|
|
|
|
def __getitem__(self, idx):
|
|
data = self.dataset[idx]
|
|
text, summary = data[self.text_column], data[self.summary_column]
|
|
text, summary = self.preprocess_fn(text, summary)
|
|
prompt = random.choice(SUMMARIZATION_SPECIAL_TOKENS["Summary"])
|
|
|
|
return (
|
|
"".join([SUMMARIZATION_SPECIAL_TOKENS["Text"], ' '.join(text.split(' ')[:256]), prompt]),
|
|
summary
|
|
)
|
|
|