Spaces:
Sleeping
Sleeping
File size: 2,235 Bytes
3d9ba5b | 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 | import json
import torch
from torch.utils.data import Dataset
from typing import List, Dict, Tuple
from nlp_core.tokenizer import BilingualTokenizer
from nlp_core.vocabulary import Vocabulary
class SummarizationDataset(Dataset):
"""
PyTorch Dataset for Sequence-to-Sequence Summarization.
"""
def __init__(
self,
data_path: str,
src_vocab: Vocabulary,
trg_vocab: Vocabulary,
lang: str = "en",
max_src_len: int = 120,
max_trg_len: int = 40
):
self.src_vocab = src_vocab
self.trg_vocab = trg_vocab
self.lang = lang
self.max_src_len = max_src_len
self.max_trg_len = max_trg_len
self.tokenizer = BilingualTokenizer()
with open(data_path, "r", encoding="utf-8") as f:
self.samples = json.load(f)
def __len__(self) -> int:
return len(self.samples)
def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
item = self.samples[idx]
article = item["article"]
summary = item["summary"]
# Tokenize
src_tokens = self.tokenizer.tokenize_words(article, lang=self.lang, remove_stopwords=False)
trg_tokens = self.tokenizer.tokenize_words(summary, lang=self.lang, remove_stopwords=False)
# Encode to indices
src_indices = self.src_vocab.encode(src_tokens, add_sos=True, add_eos=True, max_len=self.max_src_len)
trg_indices = self.trg_vocab.encode(trg_tokens, add_sos=True, add_eos=True, max_len=self.max_trg_len)
return torch.tensor(src_indices, dtype=torch.long), torch.tensor(trg_indices, dtype=torch.long)
def pad_collate_fn(batch: List[Tuple[torch.Tensor, torch.Tensor]], pad_idx: int = 0) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Collates a list of (src, trg) tensor pairs and dynamically pads them
to the max length within the mini-batch.
"""
src_list, trg_list = zip(*batch)
# Pad source tensors
src_padded = torch.nn.utils.rnn.pad_sequence(src_list, batch_first=True, padding_value=pad_idx)
# Pad target tensors
trg_padded = torch.nn.utils.rnn.pad_sequence(trg_list, batch_first=True, padding_value=pad_idx)
return src_padded, trg_padded
|