211 lines
8.8 KiB
Python
211 lines
8.8 KiB
Python
"""Restricted readers and split preparation for the official Q2 inputs."""
|
|
from __future__ import annotations
|
|
|
|
import pickle
|
|
import sys
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
|
|
ROOT = Path(__file__).resolve().parents[2]
|
|
ALIGNED_PATH = ROOT / "E题数据" / "附件2-数据集特征文件" / "aligned_50.pkl"
|
|
ATTACHMENT3_ALIGNED = ROOT / "E题数据" / "附件3-模态缺失特征样本" / "对齐版本"
|
|
MODALITIES = ("text", "audio", "vision")
|
|
EXPECTED_DIMS = {"text": 768, "audio": 74, "vision": 35}
|
|
|
|
|
|
class RestrictedUnpickler(pickle.Unpickler):
|
|
"""Allow only primitive containers and NumPy reconstruction primitives."""
|
|
|
|
_allowed = {
|
|
("builtins", name): getattr(__import__("builtins"), name)
|
|
for name in ("set", "frozenset", "slice", "complex", "bytearray")
|
|
}
|
|
_allowed.update({
|
|
("collections", "OrderedDict"): __import__("collections").OrderedDict,
|
|
("numpy", "ndarray"): np.ndarray,
|
|
("numpy", "dtype"): np.dtype,
|
|
("numpy", "asarray"): np.asarray,
|
|
("numpy.core.multiarray", "_reconstruct"): np.core.multiarray._reconstruct,
|
|
("numpy.core.multiarray", "scalar"): np.core.multiarray.scalar,
|
|
("numpy._core.multiarray", "_reconstruct"): np.core.multiarray._reconstruct,
|
|
("numpy._core.multiarray", "scalar"): np.core.multiarray.scalar,
|
|
})
|
|
|
|
def find_class(self, module: str, name: str) -> Any:
|
|
try:
|
|
return self._allowed[(module, name)]
|
|
except KeyError as exc:
|
|
raise pickle.UnpicklingError(f"blocked pickle global: {module}.{name}") from exc
|
|
|
|
|
|
def restricted_load(path: Path) -> Any:
|
|
with path.open("rb") as stream:
|
|
return RestrictedUnpickler(stream).load()
|
|
|
|
|
|
def _decode(value: Any) -> str:
|
|
if isinstance(value, bytes):
|
|
return value.decode("utf-8", errors="replace")
|
|
if isinstance(value, np.bytes_):
|
|
return bytes(value).decode("utf-8", errors="replace")
|
|
if isinstance(value, np.ndarray) and value.shape == ():
|
|
return _decode(value.item())
|
|
return str(value)
|
|
|
|
|
|
def _one_dim(value: Any, dtype: Any | None = None) -> np.ndarray:
|
|
out = np.asarray(value)
|
|
if out.ndim > 1 and out.shape[-1] == 1:
|
|
out = out.reshape(-1)
|
|
elif out.ndim > 1 and out.shape[0] == 1:
|
|
out = out.reshape(-1)
|
|
else:
|
|
out = out.reshape(-1)
|
|
return out.astype(dtype) if dtype is not None else out
|
|
|
|
|
|
@dataclass
|
|
class SplitData:
|
|
name: str
|
|
x: dict[str, np.ndarray]
|
|
mask: np.ndarray
|
|
class_y: np.ndarray | None
|
|
regression_y: np.ndarray | None
|
|
ids: list[str]
|
|
groups: np.ndarray
|
|
alignment_audit: dict[str, Any] | None = None
|
|
|
|
@property
|
|
def n(self) -> int:
|
|
return len(self.ids)
|
|
|
|
|
|
def _extract_split(
|
|
name: str, obj: dict[str, Any], with_labels: bool,
|
|
mask_override: np.ndarray | None = None,
|
|
alignment_audit: dict[str, Any] | None = None,
|
|
) -> SplitData:
|
|
raw: dict[str, np.ndarray] = {}
|
|
masks = []
|
|
for modality in MODALITIES:
|
|
arr = np.asarray(obj[modality])
|
|
if arr.ndim != 3 or arr.shape[1] != 50 or arr.shape[2] != EXPECTED_DIMS[modality]:
|
|
raise ValueError(f"{name}.{modality}: unexpected feature shape {arr.shape}")
|
|
arr = arr.astype(np.float32)
|
|
if not np.isfinite(arr).all():
|
|
raise ValueError(f"{name}.{modality}: non-finite feature values; refusing to reinterpret them as missing")
|
|
# The dataset documentation defines all-zero aligned rows as missing.
|
|
observed = np.any(arr != 0.0, axis=-1)
|
|
raw[modality] = arr
|
|
masks.append(observed)
|
|
mask = np.stack(masks, axis=-1)
|
|
if mask_override is not None:
|
|
override = np.asarray(mask_override, bool)
|
|
if override.shape != mask.shape:
|
|
raise ValueError(f"{name}: projected mask shape {override.shape} differs from {mask.shape}")
|
|
mask = override
|
|
ids = [_decode(v) for v in _one_dim(obj["id"])]
|
|
if len(ids) != len(mask):
|
|
raise ValueError(f"{name}: id count differs from feature count")
|
|
if len(set(ids)) != len(ids):
|
|
raise ValueError(f"{name}: duplicate video$_$clip primary keys")
|
|
malformed = [sample_id for sample_id in ids if "$_$" not in sample_id or not all(sample_id.split("$_$", 1))]
|
|
if malformed:
|
|
raise ValueError(f"{name}: malformed video$_$clip keys: {malformed[:5]}")
|
|
groups = np.asarray([sample_group(v) for v in ids], dtype=str)
|
|
if with_labels:
|
|
class_y = _one_dim(obj["classification_labels"], np.int64)
|
|
regression_y = _one_dim(obj["regression_labels"], np.float32)
|
|
if len(class_y) != len(ids) or len(regression_y) != len(ids):
|
|
raise ValueError(f"{name}: label count differs from feature count")
|
|
if not np.isfinite(regression_y).all() or np.any(np.abs(regression_y) > 3.0):
|
|
raise ValueError(f"{name}: regression labels must be finite and within [-3,3]")
|
|
if not np.isin(class_y, [0, 1, 2]).all():
|
|
raise ValueError(f"{name}: expected class labels in 0,1,2")
|
|
expected_class = np.where(regression_y < 0.0, 0, np.where(regression_y == 0.0, 1, 2))
|
|
mismatch = np.flatnonzero(class_y != expected_class)
|
|
if len(mismatch):
|
|
examples = [(ids[int(i)], int(class_y[i]), float(regression_y[i])) for i in mismatch[:5]]
|
|
raise ValueError(f"{name}: polarity/regression label mismatch (sample, class, score): {examples}")
|
|
else:
|
|
class_y = regression_y = None
|
|
return SplitData(name, raw, mask, class_y, regression_y, ids, groups, alignment_audit)
|
|
|
|
|
|
def sample_group(sample_id: str) -> str:
|
|
"""Official ids are video$_$clip; group on the source video only."""
|
|
return sample_id.split("$_$", 1)[0]
|
|
|
|
|
|
def load_official_splits(path: Path = ALIGNED_PATH, *, version: str = "aligned_50") -> dict[str, SplitData]:
|
|
if version not in {"aligned_50", "unaligned_50"}:
|
|
raise ValueError(f"unsupported feature version: {version}")
|
|
obj = restricted_load(path)
|
|
required = {"train", "valid", "test"}
|
|
if not isinstance(obj, dict) or not required.issubset(obj):
|
|
raise ValueError(f"{path.name} must contain train, valid, and test dictionaries")
|
|
if version == "unaligned_50":
|
|
repo_dir = Path(__file__).resolve().parents[2]
|
|
if str(repo_dir) not in sys.path:
|
|
sys.path.insert(0, str(repo_dir))
|
|
from final.adapter import adapt_official_split
|
|
|
|
splits = {}
|
|
for name in ("train", "valid", "test"):
|
|
projected, mask, audit = adapt_official_split(obj[name])
|
|
fields = {**obj[name], **projected}
|
|
splits[name] = _extract_split(name, fields, with_labels=True,
|
|
mask_override=mask, alignment_audit=audit)
|
|
else:
|
|
splits = {name: _extract_split(name, obj[name], with_labels=True) for name in ("train", "valid", "test")}
|
|
del obj
|
|
return splits
|
|
|
|
|
|
def load_attachment3_case(path: Path) -> dict[str, np.ndarray]:
|
|
obj = restricted_load(path)
|
|
case = obj.get("test", obj)
|
|
text_bert = np.asarray(case["text_bert"])
|
|
audio = np.asarray(case["audio"])
|
|
vision = np.asarray(case["vision"])
|
|
if text_bert.ndim == 3 and text_bert.shape[0] == 1:
|
|
text_bert = text_bert[0]
|
|
if text_bert.shape != (3, 50):
|
|
raise ValueError(f"{path.name}: expected text_bert (1,3,50), got {np.asarray(case['text_bert']).shape}")
|
|
result = {"input_ids": text_bert[0].astype(np.int64), "attention_mask": text_bert[1].astype(bool), "token_type_ids": text_bert[2].astype(np.int64)}
|
|
for name, arr, dim in (("audio", audio, 74), ("vision", vision, 35)):
|
|
if arr.ndim == 3 and arr.shape[0] == 1:
|
|
arr = arr[0]
|
|
if arr.shape != (50, dim):
|
|
raise ValueError(f"{path.name}: expected {name} (1,50,{dim}), got {np.asarray(case[name]).shape}")
|
|
arr = arr.astype(np.float32)
|
|
if not np.isfinite(arr).all():
|
|
raise ValueError(f"{path.name}: {name} contains non-finite features")
|
|
result[name] = arr
|
|
return result
|
|
|
|
|
|
def fit_preprocessor(train: SplitData) -> dict[str, dict[str, np.ndarray]]:
|
|
"""Fit per-dimension mean/std on observed training rows only."""
|
|
fitted: dict[str, dict[str, np.ndarray]] = {}
|
|
for j, name in enumerate(MODALITIES):
|
|
rows = train.x[name][train.mask[:, :, j]]
|
|
mean = rows.mean(axis=0, dtype=np.float64).astype(np.float32)
|
|
std = rows.std(axis=0, dtype=np.float64).astype(np.float32)
|
|
std[std < 1e-5] = 1.0
|
|
fitted[name] = {"mean": mean, "std": std}
|
|
return fitted
|
|
|
|
|
|
def transform_split(split: SplitData, fitted: dict[str, dict[str, np.ndarray]]) -> dict[str, np.ndarray]:
|
|
output = {}
|
|
for j, name in enumerate(MODALITIES):
|
|
arr = (split.x[name] - fitted[name]["mean"]) / fitted[name]["std"]
|
|
arr = np.clip(arr, -10.0, 10.0)
|
|
arr[~split.mask[:, :, j]] = 0.0
|
|
output[name] = arr.astype(np.float32)
|
|
return output
|