Spaces:
Sleeping
Sleeping
| 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 | |