src/datasets.py

from dataclasses import dataclass
from pathlib import Path

import pandas as pd
import torch
from torch.utils.data import Dataset


@dataclass
class SplitBundle:
    all_data: pd.DataFrame
    train: pd.DataFrame
    val: pd.DataFrame
    test: pd.DataFrame


def load_split_bundle(processed_dir: str | Path) -> SplitBundle:
    processed_dir = Path(processed_dir)
    return SplitBundle(
        all_data=pd.read_csv(processed_dir / "all.csv"),
        train=pd.read_csv(processed_dir / "train.csv"),
        val=pd.read_csv(processed_dir / "val.csv"),
        test=pd.read_csv(processed_dir / "test.csv"),
    )


def build_mask(node_ids: pd.Series, size: int) -> torch.Tensor:
    mask = torch.zeros(size, dtype=torch.bool)
    mask[node_ids.to_numpy()] = True
    return mask


class TitleDataset(Dataset):
    def __init__(self, dataframe: pd.DataFrame, tokenizer, max_length: int):
        self.dataframe = dataframe.reset_index(drop=True)
        self.tokenizer = tokenizer
        self.max_length = max_length

    def __len__(self) -> int:
        return len(self.dataframe)

    def __getitem__(self, index: int) -> dict:
        row = self.dataframe.iloc[index]
        encoded = self.tokenizer(
            row["title"],
            truncation=True,
            padding="max_length",
            max_length=self.max_length,
            return_tensors="pt",
        )
        item = {key: value.squeeze(0) for key, value in encoded.items()}
        item["labels"] = torch.tensor(int(row["label_id"]), dtype=torch.long)
        item["node_id"] = torch.tensor(int(row["node_id"]), dtype=torch.long)
        return item


def tokenize_dataframe(dataframe: pd.DataFrame, tokenizer, max_length: int) -> dict[str, torch.Tensor]:
    encoded = tokenizer(
        dataframe["title"].tolist(),
        truncation=True,
        padding=True,
        max_length=max_length,
        return_tensors="pt",
    )
    encoded["labels"] = torch.tensor(dataframe["label_id"].to_numpy(), dtype=torch.long)
    encoded["node_id"] = torch.tensor(dataframe["node_id"].to_numpy(), dtype=torch.long)
    return encoded