Spaces:
Running on Zero
Running on Zero
File size: 3,447 Bytes
71e6ce1 | 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 | """
dataset.py
==========
PyTorch Dataset classes for batch training/evaluation (used if you ever
retrain a model, not needed for single-question inference in predict.py,
which calls the tokenizer directly for speed/simplicity).
"""
from typing import List
import pandas as pd
import torch
from torch.utils.data import Dataset
from .config import ANSWER_MAP, OPTION_COLS
from .preprocessing import build_lstm_text
class MCQDataset(Dataset):
"""LSTM dataset (03_lstm.ipynb) — one combined text per question."""
def __init__(self, df: pd.DataFrame, tokenizer, max_len: int, is_test: bool = False):
self.df = df.reset_index(drop=True)
self.tokenizer = tokenizer
self.max_len = max_len
self.is_test = is_test
def __len__(self):
return len(self.df)
def __getitem__(self, idx):
row = self.df.iloc[idx]
text = build_lstm_text(row["prompt"], [row[c] for c in OPTION_COLS])
ids = self.tokenizer.encode(text)
ids = self.tokenizer.pad_or_truncate(ids, self.max_len)
id_tensor = torch.tensor(ids, dtype=torch.long)
if self.is_test:
return id_tensor
label = ANSWER_MAP[row["answer"]]
return id_tensor, torch.tensor(label, dtype=torch.long)
class OptionDataset(Dataset):
"""DeBERTa dataset (04_DeBERTa.ipynb) — expects a pre-expanded pairs_df
(see preprocessing.expand_pairs), one row per (question, option) pair."""
def __init__(self, pairs_df: pd.DataFrame, tokenizer, max_len: int, is_test: bool = False):
self.df = pairs_df.reset_index(drop=True)
self.tok = tokenizer
self.max_len = max_len
self.is_test = is_test
def __len__(self):
return len(self.df)
def __getitem__(self, idx):
row = self.df.iloc[idx]
enc = self.tok(
row["text"],
max_length=self.max_len,
truncation=True,
padding="max_length",
return_tensors="pt",
)
item = {
"input_ids": enc["input_ids"].squeeze(0),
"attention_mask": enc["attention_mask"].squeeze(0),
}
if "token_type_ids" in enc:
item["token_type_ids"] = enc["token_type_ids"].squeeze(0)
if not self.is_test:
item["label"] = torch.tensor(row["label"], dtype=torch.long)
return item
class MCQDatasetHF(Dataset):
"""RoBERTa dataset (05_RoBERTa.ipynb) — AutoModelForMultipleChoice format:
5 stacked (prompt, option) encodings per question."""
def __init__(self, df: pd.DataFrame, tokenizer, max_len: int, is_test: bool = False):
self.df = df.reset_index(drop=True)
self.tok = tokenizer
self.max_len = max_len
self.is_test = is_test
def __len__(self):
return len(self.df)
def __getitem__(self, idx):
row = self.df.iloc[idx]
prompt = row["prompt"]
options = [row[c] for c in OPTION_COLS]
enc = self.tok(
[prompt] * 5, options,
max_length=self.max_len, truncation=True, padding="max_length",
return_tensors="pt",
)
item = {
"input_ids": enc["input_ids"], # (5, max_len)
"attention_mask": enc["attention_mask"], # (5, max_len)
}
if not self.is_test:
item["labels"] = torch.tensor(ANSWER_MAP[row["answer"]], dtype=torch.long)
return item |