diff --git a/openspec/changes/ts-as-modality/tasks.md b/openspec/changes/ts-as-modality/tasks.md index 6788b6a..74be003 100644 --- a/openspec/changes/ts-as-modality/tasks.md +++ b/openspec/changes/ts-as-modality/tasks.md @@ -19,7 +19,7 @@ - [x] 2.2 Projector `model/projector.py`:`Linear(256→896)+LayerNorm`(验证:输出 `[B,n_patches,896]` 对齐 Qwen2.5-0.5B hidden) - [x] 2.3 多模态拼接 `model/multimodal.py`:Qwen tokenizer tokenize 文本部分,TS token 作 inputs_embeds 插入 `[属性][时间戳][TS tok][事件][问题][回答][EOS]`,统一构造 inputs_embeds+attention_mask+labels(仅回答段非 -100)(验证:8 项单测过——seq_len ≤1024、mask 全 1、回答段 label 非 -100、TS token 计入序列、无 NaN、超长左截断)。spec 澄清:训练拼回答+EOS,用 `len(tokenizer)` 构建嵌入表以容纳 special tokens。 - [x] 2.4 训练/推理封装 `model/wrapper.py`:`MultimodalTSModel` 组合 Encoder+Projector+LLM+LoRA 挂载开关,`forward` 返 loss、`generate` 返文本,支持 `freeze_llm`(阶段①)/`enable_lora(r=16)`(阶段②)(验证:端到端 forward+backward 跑通、grad 流入 Encoder/Projector、generate 出文本、阶段①峰值 1.87GB ≪5GB 设计预算)。spec 澄清:wrapper 既接收原始 batch(series+文本列表)走端到端,也兼容预拼接 inputs_embeds;encoder/projector fp32,输出投影到 LLM dtype(bf16)。 -- [ ] 2.5 Collator `data/collator.py`:JSONL→batch 张量,处理变长 C(padding+mask)与变长文本(padding+attention_mask)(验证:batch=4 张量形状一致无 NaN) +- [x] 2.5 Collator `data/collator.py`:JSONL→wrapper 原始 batch(series 张量 + attributes/timestamps/events/question/answer 文本列表),处理变长 T(padding/trunc 到 max_T)与变长 C(padding 到 max_C)、NaN→fill;变长文本 padding/attention_mask 由 wrapper 逐样本 splice+max_seq 填充处理(验证:11 项单测过;端到端 Collator→Wrapper fwd+bwd+generate 跑通) - [ ] 2.6 M2 出口验证:单 batch 前向+反向在 3060 跑通,loss 有限且下降趋势,阶段①模式峰值显存 ~5GB ## 3. M3 · 阶段①对齐训练 diff --git a/src/tsmm/data/collator.py b/src/tsmm/data/collator.py new file mode 100644 index 0000000..0dff93f --- /dev/null +++ b/src/tsmm/data/collator.py @@ -0,0 +1,103 @@ +"""Collator (T2.5): JSONL sample dicts → wrapper raw batch. + +Input : a list of sample dicts (the gen_synthetic / real_bench JSONL schema): + { + "series": [[T, C] floats], # may contain NaN (missing) + "attributes": [{"name","unit","freq"}], + "timestamps": [T ints], + "events": [{"t","text"}], + "instruction": str, + "answer": str | None, + "labels": [T ints], + } + +Output : the wrapper's raw batch contract: + { + "series": [B, max_T, max_C] float tensor (NaN→nan_fill, T pad/trunc), + "attributes": list[str], # descriptive channel text per sample + "timestamps": list[str], # range text per sample + "events": list[str], # event text per sample + "question": list[str], # from sample["instruction"] + "answer": list[str] | None, + } + +Variable-length text (token padding / attention_mask) is handled downstream by +the wrapper's per-sample splice + max_seq padding, so the collator only needs to +produce the per-sample strings. + +Design ref: §6.1 (``data/collator.py``). +""" +from __future__ import annotations + +from typing import Any, List, Optional + +import torch + + +def _attrs_to_text(attrs: List[dict]) -> str: + parts = [] + for a in attrs: + parts.append(f"{a.get('name','?')}({a.get('unit','')},{a.get('freq','')})") + return " ".join(parts) + + +def _events_to_text(events: List[dict]) -> str: + if not events: + return "" + parts = [f"t={e.get('t')}: {e.get('text','')}" for e in events] + return "; ".join(parts) + + +def _timestamps_to_text(ts: List[int]) -> str: + if not ts: + return "" + if len(ts) == 1: + return f"t={ts[0]}" + return f"t={ts[0]}..{ts[-1]} ({len(ts)} steps)" + + +class Collator: + def __init__(self, max_T: int = 512, nan_fill: float = 0.0) -> None: + self.max_T = max_T + self.nan_fill = nan_fill + + def _stack_series(self, samples: List[dict]) -> torch.Tensor: + tensors = [] + for s in samples: + arr = torch.tensor(s["series"], dtype=torch.float32) # [T, C] + T, C = arr.shape + # truncate / pad T to max_T + if T >= self.max_T: + arr = arr[: self.max_T] + else: + pad = torch.zeros(self.max_T - T, C) + arr = torch.cat([arr, pad], dim=0) + tensors.append((arr, C)) + max_C = max(c for _, c in tensors) + B = len(tensors) + out = torch.zeros(B, self.max_T, max_C, dtype=torch.float32) + for i, (arr, C) in enumerate(tensors): + out[i, :, :C] = arr[:, :max_C] + # replace NaN (missing markers) with fill + out = torch.nan_to_num(out, nan=self.nan_fill) + return out + + def __call__(self, samples: List[dict]) -> dict: + series = self._stack_series(samples) + attributes = [_attrs_to_text(s.get("attributes", [])) for s in samples] + timestamps = [_timestamps_to_text(s.get("timestamps", [])) for s in samples] + events = [_events_to_text(s.get("events", [])) for s in samples] + question = [s.get("instruction", "") for s in samples] + answers = [s.get("answer") for s in samples] + if all(a is None for a in answers): + answer: Optional[List[str]] = None + else: + answer = ["" if a is None else a for a in answers] + return { + "series": series, + "attributes": attributes, + "timestamps": timestamps, + "events": events, + "question": question, + "answer": answer, + } diff --git a/tests/test_collator.py b/tests/test_collator.py new file mode 100644 index 0000000..032b64f --- /dev/null +++ b/tests/test_collator.py @@ -0,0 +1,97 @@ +"""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