File size: 2,106 Bytes
877049d | 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 | 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
|