| """ |
| PyTorch Datasets and DataLoaders for Document Pre-training and QA Fine-Tuning. |
| """ |
| import torch |
| from torch.utils.data import Dataset |
| from typing import List, Dict, Any, Optional, Tuple |
| from tokenizer.tokenizer import CustomBPETokenizer |
|
|
|
|
| class DocPretrainDataset(Dataset): |
| """ |
| Dataset for next-token autoregressive training on extracted document text. |
| Segments token stream into fixed-length context windows with stride. |
| """ |
| def __init__( |
| self, |
| text_corpus: List[str], |
| tokenizer: CustomBPETokenizer, |
| seq_len: int = 256, |
| stride: int = 128 |
| ): |
| self.tokenizer = tokenizer |
| self.seq_len = seq_len |
| self.samples = [] |
|
|
| |
| full_tokens = [] |
| for text in text_corpus: |
| tokens = tokenizer.encode(text, add_bos=True, add_eos=True) |
| full_tokens.extend(tokens) |
|
|
| |
| if len(full_tokens) <= seq_len + 1: |
| |
| pad_len = (seq_len + 1) - len(full_tokens) |
| sample = full_tokens + [tokenizer.pad_id] * pad_len |
| self.samples.append(sample) |
| else: |
| for i in range(0, len(full_tokens) - seq_len, stride): |
| chunk = full_tokens[i : i + seq_len + 1] |
| if len(chunk) == seq_len + 1: |
| self.samples.append(chunk) |
|
|
| def __len__(self) -> int: |
| return len(self.samples) |
|
|
| def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]: |
| tokens = self.samples[idx] |
| x = torch.tensor(tokens[:-1], dtype=torch.long) |
| y = torch.tensor(tokens[1:], dtype=torch.long) |
| return x, y |
|
|
|
|
| class DocQADataset(Dataset): |
| """ |
| Dataset for Context-conditioned Document Question-Answering. |
| Formats samples as: |
| <bos> <doc> {context} <sep> <query> {question} <sep> <ans> {answer} <eos> |
| """ |
| def __init__( |
| self, |
| qa_pairs: List[Dict[str, str]], |
| tokenizer: CustomBPETokenizer, |
| max_seq_len: int = 512, |
| mask_prompt_loss: bool = True |
| ): |
| """ |
| qa_pairs: List of {"context": str, "question": str, "answer": str} |
| mask_prompt_loss: If True, only compute gradient loss on the answer tokens. |
| """ |
| self.tokenizer = tokenizer |
| self.max_seq_len = max_seq_len |
| self.mask_prompt_loss = mask_prompt_loss |
| self.samples = [] |
|
|
| for item in qa_pairs: |
| context = item.get("context", "") |
| question = item.get("question", "") |
| answer = item.get("answer", "") |
|
|
| |
| prompt_str = f"<bos> <doc> {context} <sep> <query> {question} <sep> <ans> " |
| answer_str = f"{answer} <eos>" |
|
|
| prompt_ids = tokenizer.encode(prompt_str, add_bos=False, add_eos=False) |
| answer_ids = tokenizer.encode(answer_str, add_bos=False, add_eos=False) |
|
|
| total_ids = prompt_ids + answer_ids |
|
|
| if len(total_ids) > max_seq_len: |
| |
| max_ctx_len = max_seq_len - len(answer_ids) - 10 |
| if max_ctx_len > 10: |
| truncated_prompt_str = f"<bos> <doc> {context[:max_ctx_len*3]} <sep> <query> {question} <sep> <ans> " |
| prompt_ids = tokenizer.encode(truncated_prompt_str, add_bos=False, add_eos=False) |
| total_ids = (prompt_ids + answer_ids)[:max_seq_len] |
| else: |
| total_ids = total_ids[:max_seq_len] |
|
|
| self.samples.append({ |
| "input_ids": total_ids[:-1], |
| "target_ids": total_ids[1:], |
| "prompt_len": len(prompt_ids) - 1 |
| }) |
|
|
| def __len__(self) -> int: |
| return len(self.samples) |
|
|
| def __getitem__(self, idx: int) -> Dict[str, Any]: |
| item = self.samples[idx] |
| x = list(item["input_ids"]) |
| y = list(item["target_ids"]) |
| prompt_len = item["prompt_len"] |
|
|
| |
| if self.mask_prompt_loss: |
| |
| for p in range(min(prompt_len, len(y))): |
| y[p] = -1 |
|
|
| return { |
| "input_ids": torch.tensor(x, dtype=torch.long), |
| "target_ids": torch.tensor(y, dtype=torch.long), |
| "length": len(x) |
| } |
|
|
|
|
| def qa_collate_fn(batch: List[Dict[str, Any]], pad_id: int = 0) -> Tuple[torch.Tensor, torch.Tensor]: |
| """ |
| Collate variable-length sequences with right-padding. |
| """ |
| max_len = max(item["length"] for item in batch) |
| batch_x = [] |
| batch_y = [] |
|
|
| for item in batch: |
| x = item["input_ids"] |
| y = item["target_ids"] |
| pad_len = max_len - len(x) |
|
|
| if pad_len > 0: |
| padded_x = torch.cat([x, torch.full((pad_len,), pad_id, dtype=torch.long)]) |
| padded_y = torch.cat([y, torch.full((pad_len,), -1, dtype=torch.long)]) |
| else: |
| padded_x = x |
| padded_y = y |
|
|
| batch_x.append(padded_x) |
| batch_y.append(padded_y) |
|
|
| return torch.stack(batch_x), torch.stack(batch_y) |
|
|
