Spaces:
Runtime error
Runtime error
File size: 4,742 Bytes
fa6856c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 |
import json
import pandas as pd
import torch
from datasets import load_dataset
from torch.utils.data import Dataset
def get_dataset_from_jsonl(jsonl_file, return_summary=True):
# if return_summary is True, return a list of posts with summary concatenated
# if return_summary is False, return a list of posts and a list of summaries
with open(jsonl_file, "r") as f:
dataset = [json.loads(line) for line in f]
post_list = []
summary_list = []
for d in dataset:
if return_summary:
post = f"SUBREDDIT: r/{d['subreddit']}\nTITLE: {d['title']}\nPOST: {d['post']}\nTL;DR: {d['summary']}"
else:
post = f"SUBREDDIT: r/{d['subreddit']}\nTITLE: {d['title']}\nPOST: {d['post']}\nTL;DR: "
summary_list.append(d["summary"])
post_list.append(post)
if not return_summary:
return post_list, summary_list
return post_list
class TLDRDataset(Dataset):
def __init__(self, train_path, tokenizer, split, max_length=550):
self.post_list = []
dataset = load_dataset(train_path, split=split)
for sample in dataset:
self.post_list.append(sample["prompt"] + sample["label"])
if "valid" in split:
self.post_list = self.post_list[0:2000]
self.tokenizer = tokenizer
self.max_length = max_length
self.input_ids = []
self.attn_masks = []
def __len__(self):
return len(self.post_list)
def __getitem__(self, idx):
txt = self.post_list[idx]
encodings_dict = self.tokenizer(txt, truncation=True, max_length=self.max_length, padding="max_length")
input_ids = torch.tensor(encodings_dict["input_ids"])
attn_masks = torch.tensor(encodings_dict["attention_mask"])
return {
"input_ids": input_ids,
"attention_mask": attn_masks,
"labels": input_ids,
}
class ComparisonDataset(Dataset):
def __init__(self, comparison_path, tokenizer, max_length=550):
with open(comparison_path, "r") as f:
dataset = [json.loads(line) for line in f]
self.tokenizer = tokenizer
self.post_list = []
self.summaries_0 = []
self.summaries_1 = []
self.labels = []
self.max_length = max_length
def make_text(post, summarize):
return f"SUBREDDIT: r/{post['subreddit']}\nTITLE: {post['title']}\nPOST: {post['post']}\nTL;DR: {summarize}"
for sample in dataset: # chosen summary is always the first one
self.post_list.append(sample["info"]["post"])
# NOTE: The chosen summary is always the first one, i.e. `sample["summaries"][0]`
if sample["choice"] == 0:
self.summaries_0.append(make_text(sample["info"], sample["summaries"][0]["text"]))
self.summaries_1.append(make_text(sample["info"], sample["summaries"][1]["text"]))
else:
self.summaries_0.append(make_text(sample["info"], sample["summaries"][1]["text"]))
self.summaries_1.append(make_text(sample["info"], sample["summaries"][0]["text"]))
self.labels.append(0)
def __len__(self):
return len(self.post_list)
def __getitem__(self, idx):
summ0 = self.summaries_0[idx]
summ1 = self.summaries_1[idx]
encodings_dict = self.tokenizer(
[summ0, summ1],
truncation=True,
max_length=self.max_length,
padding="max_length",
)
input_ids = torch.tensor(encodings_dict["input_ids"])
attention_mask = torch.tensor(encodings_dict["attention_mask"])
return {"input_ids": input_ids, "attention_mask": attention_mask}
class AllSummDataset(Dataset):
def __init__(self, train_path, tokenizer, split, max_length=1024):
df = pd.read_parquet(train_path)
if split == "valid":
df = df.sample(n=5000)
self.summarizes = []
for i, row in df.iterrows():
self.summarizes.append(f"Summarize: {row['text']}. TL;DR: {row['summary']}")
self.tokenizer = tokenizer
self.max_length = max_length
self.input_ids = []
self.attn_masks = []
def __len__(self):
return len(self.summarizes)
def __getitem__(self, idx):
txt = self.summarizes[idx]
encodings_dict = self.tokenizer(txt, truncation=True, max_length=self.max_length, padding="max_length")
input_ids = torch.tensor(encodings_dict["input_ids"])
attn_masks = torch.tensor(encodings_dict["attention_mask"])
return {
"input_ids": input_ids,
"attention_mask": attn_masks,
"labels": input_ids,
}
|