File size: 4,408 Bytes
38136ed | 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 109 110 111 112 113 114 115 116 | """PyTorch Dataset classes.
We provide three:
- ACSADataset text only, per-aspect labels -> Baseline 3 (no meta)
- MetaACSADataset text + meta features -> Proposed model
- OverallSentimentDataset text only, single 3-class -> Baseline 2 (BERT-overall)
"""
from typing import Optional
import numpy as np
import pandas as pd
import torch
from torch.utils.data import Dataset
from . import config as cfg
from .meta_encoder import MetaEncoder
class _BaseTextDataset(Dataset):
def __init__(self, df: pd.DataFrame, tokenizer, max_length: int = cfg.MAX_LENGTH,
text_col: str = "full_text"):
self.df = df.reset_index(drop=True)
self.tokenizer = tokenizer
self.max_length = max_length
self.text_col = text_col
def __len__(self):
return len(self.df)
def _encode_text(self, text: str):
enc = self.tokenizer(
text,
max_length=self.max_length,
padding="max_length",
truncation=True,
return_tensors="pt",
)
return enc["input_ids"].squeeze(0), enc["attention_mask"].squeeze(0)
class ACSADataset(_BaseTextDataset):
"""Per-aspect multi-head training, text-only. (Baseline 3.)"""
def __init__(self, df, tokenizer, **kw):
super().__init__(df, tokenizer, **kw)
self.aspect_cols = [f"aspect_{a}" for a in cfg.ASPECTS]
missing = [c for c in self.aspect_cols if c not in self.df.columns]
if missing:
raise KeyError(f"ACSADataset: missing aspect label columns {missing}. "
f"Did you run weak labeling (script 03)?")
def __getitem__(self, idx):
row = self.df.iloc[idx]
input_ids, attn = self._encode_text(str(row[self.text_col]))
labels = torch.tensor([int(row[c]) for c in self.aspect_cols], dtype=torch.long)
return {"input_ids": input_ids, "attention_mask": attn, "labels": labels}
class MetaACSADataset(_BaseTextDataset):
"""Per-aspect multi-head training, text + metadata. (Proposed model.)
The metadata is pre-encoded to a numpy matrix of shape (N, meta_dim) using
a fit MetaEncoder, and indexed in lock-step with the text rows.
If the df contains an 'overall_label' column (0=Neg, 1=Neu, 2=Pos), it is
returned as part of each batch for the joint overall auxiliary loss.
"""
def __init__(self, df, tokenizer, meta_encoder: MetaEncoder, **kw):
super().__init__(df, tokenizer, **kw)
self.aspect_cols = [f"aspect_{a}" for a in cfg.ASPECTS]
missing = [c for c in self.aspect_cols if c not in self.df.columns]
if missing:
raise KeyError(f"MetaACSADataset: missing aspect label columns {missing}.")
self.has_overall = "overall_label" in self.df.columns
# Pre-encode metadata once, in the same row order as self.df.
self.meta_matrix = meta_encoder.transform(self.df)
assert self.meta_matrix.shape[0] == len(self.df), \
"meta matrix and df length mismatch"
def __getitem__(self, idx):
row = self.df.iloc[idx]
input_ids, attn = self._encode_text(str(row[self.text_col]))
labels = torch.tensor([int(row[c]) for c in self.aspect_cols], dtype=torch.long)
meta = torch.from_numpy(self.meta_matrix[idx]).float()
item = {
"input_ids": input_ids,
"attention_mask": attn,
"meta_features": meta,
"labels": labels,
}
if self.has_overall:
item["overall_labels"] = torch.tensor(int(row["overall_label"]), dtype=torch.long)
return item
class OverallSentimentDataset(_BaseTextDataset):
"""Single 3-class overall sentiment label. (Baseline 2: BERT-overall.)"""
def __init__(self, df, tokenizer, label_col: str = "overall_label", **kw):
super().__init__(df, tokenizer, **kw)
self.label_col = label_col
if self.label_col not in self.df.columns:
raise KeyError(f"OverallSentimentDataset: missing {self.label_col} column.")
def __getitem__(self, idx):
row = self.df.iloc[idx]
input_ids, attn = self._encode_text(str(row[self.text_col]))
return {
"input_ids": input_ids,
"attention_mask": attn,
"labels": torch.tensor(int(row[self.label_col]), dtype=torch.long),
}
|