张宗平
7e1df76e3e
- src/tsmm/data/collator.py: Collator maps JSONL sample dicts to the wrapper
batch contract {series [B,max_T,max_C], attributes, timestamps, events,
question, answer}. Variable-T pad/truncate, variable-C pad, NaN→fill.
Variable-text padding/attention_mask deferred to the wrapper's per-sample
splice + max_seq padding.
- tests/test_collator.py: 11 tests. End-to-end Collator→Wrapper fwd+bwd+generate
verified (loss finite, peak 2.32 GB, generate emits text). 91 tests passing.
98 lines
3.5 KiB
Python
98 lines
3.5 KiB
Python
"""Tests for Collator (T2.5): JSONL sample dicts → wrapper raw batch."""
|
|
import math
|
|
|
|
import torch
|
|
|
|
from tsmm.data.collator import Collator
|
|
|
|
|
|
def make_sample(T=8, C=3, with_nan=False, n_events=1, answer="no", instruction="Is there an anomaly?"):
|
|
"""Build a dict mimicking the gen_pipeline JSONL schema."""
|
|
series = [[float(i * C + c) for c in range(C)] for i in range(T)]
|
|
if with_nan:
|
|
series[1][0] = float("nan")
|
|
return {
|
|
"series": series, # [T, C]
|
|
"attributes": [
|
|
{"name": f"var{c}", "unit": "%", "freq": "1Hz"} for c in range(C)
|
|
],
|
|
"timestamps": [1700000000 + i for i in range(T)],
|
|
"events": [{"t": 3, "text": "deploy"}] * n_events,
|
|
"instruction": instruction,
|
|
"answer": answer,
|
|
"labels": [0] * T,
|
|
}
|
|
|
|
|
|
class TestCollator:
|
|
def test_basic_shapes_and_types(self):
|
|
col = Collator(max_T=8)
|
|
batch = col([make_sample(T=8, C=3) for _ in range(4)])
|
|
assert batch["series"].shape == (4, 8, 3)
|
|
assert isinstance(batch["attributes"], list) and len(batch["attributes"]) == 4
|
|
assert isinstance(batch["timestamps"], list) and len(batch["timestamps"]) == 4
|
|
assert isinstance(batch["events"], list) and len(batch["events"]) == 4
|
|
assert isinstance(batch["question"], list) and len(batch["question"]) == 4
|
|
assert isinstance(batch["answer"], list) and len(batch["answer"]) == 4
|
|
|
|
def test_question_maps_from_instruction(self):
|
|
col = Collator()
|
|
batch = col([make_sample()])
|
|
assert batch["question"][0] == "Is there an anomaly?"
|
|
|
|
def test_answer_passthrough(self):
|
|
col = Collator()
|
|
batch = col([make_sample(answer="yes")])
|
|
assert batch["answer"][0] == "yes"
|
|
|
|
def test_no_nan_after_collate(self):
|
|
col = Collator(nan_fill=0.0)
|
|
batch = col([make_sample(T=8, C=3, with_nan=True)])
|
|
assert not torch.isnan(batch["series"]).any()
|
|
# the NaN position was filled with 0
|
|
assert batch["series"][0, 1, 0].item() == 0.0
|
|
|
|
def test_truncates_long_T(self):
|
|
col = Collator(max_T=16)
|
|
batch = col([make_sample(T=32, C=2)])
|
|
assert batch["series"].shape == (1, 16, 2)
|
|
|
|
def test_pads_short_T(self):
|
|
col = Collator(max_T=64)
|
|
batch = col([make_sample(T=16, C=2)])
|
|
assert batch["series"].shape == (1, 64, 2)
|
|
# padded region is zero
|
|
assert batch["series"][0, 16:].abs().sum().item() == 0.0
|
|
|
|
def test_pads_variable_C_within_batch(self):
|
|
col = Collator(max_T=16)
|
|
batch = col([make_sample(T=16, C=2), make_sample(T=16, C=5)])
|
|
# both padded to max_C=5
|
|
assert batch["series"].shape == (2, 16, 5)
|
|
|
|
def test_attribute_string_is_descriptive(self):
|
|
col = Collator()
|
|
batch = col([make_sample(C=2)])
|
|
s = batch["attributes"][0]
|
|
assert "var0" in s and "var1" in s
|
|
|
|
def test_timestamps_string(self):
|
|
col = Collator()
|
|
batch = col([make_sample(T=8)])
|
|
s = batch["timestamps"][0]
|
|
assert isinstance(s, str) and "1700000000" in s
|
|
|
|
def test_events_string(self):
|
|
col = Collator()
|
|
batch = col([make_sample(n_events=2)])
|
|
s = batch["events"][0]
|
|
assert isinstance(s, str) and "deploy" in s
|
|
|
|
def test_answer_optional(self):
|
|
# answer=None path (inference): collator should omit/None answer
|
|
col = Collator()
|
|
sample = make_sample()
|
|
sample["answer"] = None
|
|
batch = col([sample])
|
|
assert batch.get("answer") is None or batch["answer"][0] is None
|