File size: 1,808 Bytes
4d67ec7 | 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 | """
dataset_multilabel.py — PyTorch Dataset for Multi-Label NEET Biology classification
Expected JSON format:
[
{
"question" : "Which structures are common to both plant cells and animal cells?",
"chapter_ids" : [7, 8] # list of ints
},
...
]
"""
import json
import torch
from torch.utils.data import Dataset
from transformers import AutoTokenizer
class NEETMultiLabelDataset(Dataset):
def __init__(self, data_path: str, tokenizer: AutoTokenizer, num_labels: int, max_length: int = 128):
with open(data_path, "r", encoding="utf-8") as f:
self.data = json.load(f)
self.tokenizer = tokenizer
self.max_length = max_length
self.num_labels = num_labels
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
item = self.data[idx]
question = item["question"]
# This will handle both string dicts or lists of IDs depending on generation
# We expect item["chapter_ids"] to be a list of integers
label_ids = item.get("chapter_ids", [])
if not isinstance(label_ids, list):
label_ids = [label_ids]
# Create a one-hot float vector for BCEWithLogitsLoss
labels = torch.zeros(self.num_labels, dtype=torch.float)
for c_id in label_ids:
if 0 <= c_id < self.num_labels:
labels[c_id] = 1.0
encoding = self.tokenizer(
question,
max_length=self.max_length,
padding="max_length",
truncation=True,
return_tensors="pt",
)
return {
"input_ids": encoding["input_ids"].squeeze(0),
"attention_mask": encoding["attention_mask"].squeeze(0),
"labels": labels,
}
|