Spaces:
Running
Running
File size: 7,455 Bytes
590a604 ee1a8a3 1fbc47b 2286a5e 1fbc47b 2286a5e 1fbc47b 2286a5e 1fbc47b ee1a8a3 1fbc47b |
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 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 |
"""
Dataset definitions for the LexiMind multitask training pipeline.
Defines PyTorch Dataset classes and data loading utilities for summarization,
emotion classification, and topic classification tasks. Supports both JSON
array and JSONL file formats.
Author: Oliver Perrin
Date: December 2025
"""
from __future__ import annotations
import json
from dataclasses import dataclass
from pathlib import Path
from typing import Callable, Iterable, List, Sequence, TypeVar
from sklearn.preprocessing import LabelEncoder, MultiLabelBinarizer
from torch.utils.data import Dataset
@dataclass
class SummarizationExample:
"""Container for abstractive summarization samples."""
source: str
summary: str
@dataclass
class EmotionExample:
"""Container for multi-label emotion classification samples."""
text: str
emotions: Sequence[str]
@dataclass
class TopicExample:
"""Container for topic clustering / classification samples."""
text: str
topic: str
class SummarizationDataset(Dataset[SummarizationExample]):
"""Dataset yielding encoder-decoder training pairs."""
def __init__(self, examples: Iterable[SummarizationExample]) -> None:
self._examples = list(examples)
def __len__(self) -> int:
return len(self._examples)
def __getitem__(self, index: int) -> SummarizationExample:
return self._examples[index]
class EmotionDataset(Dataset[EmotionExample]):
"""Dataset that owns a scikit-learn MultiLabelBinarizer for emissions."""
def __init__(
self,
examples: Iterable[EmotionExample],
*,
binarizer: MultiLabelBinarizer | None = None,
) -> None:
self._examples = list(examples)
all_labels = [example.emotions for example in self._examples]
if binarizer is None:
self._binarizer = MultiLabelBinarizer()
self._binarizer.fit(all_labels)
else:
self._binarizer = binarizer
if not hasattr(self._binarizer, "classes_"):
raise ValueError(
"Provided MultiLabelBinarizer must be pre-fitted with 'classes_' attribute."
)
def __len__(self) -> int:
return len(self._examples)
def __getitem__(self, index: int) -> EmotionExample:
return self._examples[index]
@property
def binarizer(self) -> MultiLabelBinarizer:
return self._binarizer
@property
def emotion_classes(self) -> List[str]:
return list(self._binarizer.classes_)
class TopicDataset(Dataset[TopicExample]):
"""Dataset that owns a LabelEncoder for topic ids."""
def __init__(
self,
examples: Iterable[TopicExample],
*,
encoder: LabelEncoder | None = None,
) -> None:
self._examples = list(examples)
topics = [example.topic for example in self._examples]
if encoder is None:
self._encoder = LabelEncoder().fit(topics)
else:
self._encoder = encoder
if not hasattr(self._encoder, "classes_"):
raise ValueError(
"Provided LabelEncoder must be pre-fitted with 'classes_' attribute."
)
def __len__(self) -> int:
return len(self._examples)
def __getitem__(self, index: int) -> TopicExample:
return self._examples[index]
@property
def encoder(self) -> LabelEncoder:
return self._encoder
@property
def topic_classes(self) -> List[str]:
return list(self._encoder.classes_)
T = TypeVar("T")
def _safe_json_load(handle, path: Path) -> object:
try:
return json.load(handle)
except json.JSONDecodeError as exc:
raise ValueError(f"Failed to parse JSON in '{path}': {exc}") from exc
def _safe_json_loads(data: str, path: Path, line_number: int) -> object:
try:
return json.loads(data)
except json.JSONDecodeError as exc:
raise ValueError(f"Failed to parse JSON in '{path}' at line {line_number}: {exc}") from exc
def _validate_keys(
payload: dict,
required_keys: Sequence[str],
position: int,
*,
path: Path,
is_array: bool = False,
) -> None:
missing = [key for key in required_keys if key not in payload]
if missing:
keys = ", ".join(sorted(missing))
location = "index" if is_array else "line"
raise KeyError(f"Missing required keys ({keys}) at {location} {position} of '{path}'")
def _load_jsonl_generic(
path: str,
constructor: Callable[[dict], T],
required_keys: Sequence[str],
) -> List[T]:
data_path = Path(path)
if not data_path.exists():
raise FileNotFoundError(f"Dataset file '{data_path}' does not exist")
if not data_path.is_file():
raise ValueError(f"Dataset path '{data_path}' is not a file")
items: List[T] = []
with data_path.open("r", encoding="utf-8") as handle:
first_non_ws = ""
while True:
pos = handle.tell()
char = handle.read(1)
if not char:
break
if not char.isspace():
first_non_ws = char
handle.seek(pos)
break
if not first_non_ws:
raise ValueError(f"Dataset file '{data_path}' is empty or contains only whitespace")
if first_non_ws == "[":
payloads = _safe_json_load(handle, data_path)
if not isinstance(payloads, list):
raise ValueError(
f"Expected a JSON array in '{data_path}' but found {type(payloads).__name__}"
)
for idx, payload in enumerate(payloads):
if not isinstance(payload, dict):
raise ValueError(
f"Expected objects in array for '{data_path}', found {type(payload).__name__} at index {idx}"
)
_validate_keys(payload, required_keys, idx, path=data_path, is_array=True)
items.append(constructor(payload))
else:
handle.seek(0)
line_number = 0
for line in handle:
line_number += 1
if not line.strip():
continue
payload = _safe_json_loads(line, data_path, line_number)
if not isinstance(payload, dict):
raise ValueError(
f"Expected JSON object per line in '{data_path}', found {type(payload).__name__} at line {line_number}"
)
_validate_keys(payload, required_keys, line_number, path=data_path)
items.append(constructor(payload))
return items
def load_summarization_jsonl(path: str) -> List[SummarizationExample]:
return _load_jsonl_generic(
path,
lambda payload: SummarizationExample(source=payload["source"], summary=payload["summary"]),
required_keys=("source", "summary"),
)
def load_emotion_jsonl(path: str) -> List[EmotionExample]:
return _load_jsonl_generic(
path,
lambda payload: EmotionExample(text=payload["text"], emotions=payload.get("emotions", [])),
required_keys=("text",),
)
def load_topic_jsonl(path: str) -> List[TopicExample]:
return _load_jsonl_generic(
path,
lambda payload: TopicExample(text=payload["text"], topic=payload["topic"]),
required_keys=("text", "topic"),
)
|