Flatten submit package structure

submit/q2/math/__init__.py

@@ -0,0 +1 @@
+"""Mathematical Q2 model family and training driver."""

submit/q2/math/data.py

@@ -0,0 +1,213 @@
+"""Restricted readers and split preparation for the official Q2 inputs."""
+from __future__ import annotations
+
+import pickle
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+
+from data_paths import ATTACHMENT2, ATTACHMENT3, DATA_ROOT, PROJECT_ROOT
+
+ROOT = PROJECT_ROOT
+ATTACHMENT2_DIR = ATTACHMENT2
+ALIGNED_PATH = ATTACHMENT2 / "aligned_50.pkl"
+ATTACHMENT3_ALIGNED = ATTACHMENT3 / "对齐版本"
+ATTACHMENT3_UNALIGNED = ATTACHMENT3 / "未对齐版本"
+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,
+    })
+    if hasattr(np.core.numeric, "_frombuffer"):
+        _allowed[("numpy.core.numeric", "_frombuffer")] = np.core.numeric._frombuffer
+        _allowed[("numpy._core.numeric", "_frombuffer")] = np.core.numeric._frombuffer
+
+    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":
+        from 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

submit/q2/math/predict_attachment3.py

@@ -0,0 +1,99 @@
+"""Run a saved Q2 model on the unlabeled Attachment 3 cases."""
+from __future__ import annotations
+
+import argparse
+import json
+import time
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from data_paths import PROJECT_ROOT
+from model.crg import INPUT_DIMS, MODALITIES, StructuredGaussianImputer
+from .train import (
+    _make_variant,
+    infer_attachment3,
+    reencode_attachment3,
+    validate_attachment3_predictions,
+    write_csv,
+)
+
+DEFAULT_RESULTS_DIR = PROJECT_ROOT / "experiments" / "q2" / "unaligned_math_all_b128"
+DEFAULT_OUTPUT_DIR = PROJECT_ROOT / "output" / "q2"
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--input-version", choices=("aligned_50", "unaligned_50"), default="unaligned_50")
+    parser.add_argument("--results-dir", type=Path, default=DEFAULT_RESULTS_DIR,
+                        help="saved Q2 checkpoint and calibration directory")
+    parser.add_argument("--output-dir", type=Path, default=DEFAULT_OUTPUT_DIR)
+    parser.add_argument("--device", choices=("auto", "cpu", "cuda"), default="auto")
+    args = parser.parse_args()
+
+    if args.device == "cuda" and not torch.cuda.is_available():
+        parser.error("CUDA was requested but is not available")
+    device_name = "cuda" if args.device == "auto" and torch.cuda.is_available() else args.device
+    if device_name == "auto":
+        device_name = "cpu"
+    device = torch.device(device_name)
+
+    results_dir = args.results_dir.expanduser().resolve()
+    output_dir = args.output_dir.expanduser().resolve()
+    manifest_path = results_dir / "run_manifest.json"
+    calibration_path = results_dir / "validation_metrics.json"
+    manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
+    calibration = json.loads(calibration_path.read_text(encoding="utf-8"))
+    selected = calibration.get("selected_model", manifest.get("selected_model"))
+    if not selected:
+        raise ValueError(f"no selected_model recorded in {calibration_path}")
+
+    imputer = StructuredGaussianImputer(INPUT_DIMS).to(device)
+    imputer.load_state_dict(torch.load(results_dir / "structured_imputer.pt", map_location=device, weights_only=True))
+    model = _make_variant(selected, imputer).to(device)
+    model.load_state_dict(torch.load(results_dir / "crg_student.pt", map_location=device, weights_only=True))
+
+    with np.load(results_dir / "preprocessor.npz", allow_pickle=False) as archive:
+        fitted = {m: {k: archive[f"{m}_{k}"].copy() for k in ("mean", "std")} for m in MODALITIES}
+    priors = manifest["attachment3_low_information_priors"]
+    temperature = float(calibration["temperature"])
+    class_prior = np.asarray(priors["class_probability_values"], dtype=np.float64)
+    magnitude_priors = np.asarray((priors["negative_beta"], priors["positive_beta"]), dtype=np.float32)
+
+    cases, source_audit = reencode_attachment3(device, input_version=args.input_version)
+    predictions, inference_audit = infer_attachment3(
+        model, cases, fitted, device, temperature, class_prior, magnitude_priors,
+    )
+    validate_attachment3_predictions([case["case_id"] for case in cases], predictions)
+
+    output_dir.mkdir(parents=True, exist_ok=True)
+    inference_by_id = {row["case_id"]: row for row in inference_audit}
+    predictions_path = output_dir / "attachment3_predictions.csv"
+    audit_path = output_dir / "attachment3_audit.csv"
+    manifest_out_path = output_dir / "attachment3_prediction_manifest.json"
+    write_csv(predictions_path, predictions)
+    write_csv(audit_path, [{**source, **inference_by_id[source["case_id"]]} for source in source_audit])
+
+    try:
+        results_reference = results_dir.relative_to(PROJECT_ROOT).as_posix()
+    except ValueError:
+        results_reference = "external checkpoint directory"
+    prediction_manifest = {
+        "task": "unlabeled Attachment 3 inference",
+        "input_version": args.input_version,
+        "selected_model": selected,
+        "checkpoint_run": results_reference,
+        "prediction_count": len(predictions),
+        "temperature": temperature,
+        "labels_available": False,
+        "prediction_file": predictions_path.name,
+        "audit_file": audit_path.name,
+        "completed_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
+    }
+    manifest_out_path.write_text(json.dumps(prediction_manifest, ensure_ascii=False, indent=2), encoding="utf-8")
+    print(f"Wrote {len(predictions)} unlabeled Attachment 3 predictions to {output_dir}", flush=True)
+
+
+if __name__ == "__main__":
+    main()