from __future__ import annotations
import random
import subprocess
from collections import OrderedDict
from collections.abc import Sequence
from pathlib import Path
from typing import Any
import torch
from torch.nn.utils.rnn import pad_sequence
from torch.utils.data import DistributedSampler, Sampler
TS_START_TOKEN = ""
TS_END_TOKEN = ""
SCALE_START_TOKEN = ""
SCALE_END_TOKEN = ""
PromptMode = str
STAGE1_PROMPT_VARIANTS_UNIVAR: tuple[str, ...] = (
"请描述这个时间序列：{ts_start} {ts_end}",
"请概括这个时间序列：{ts_start} {ts_end}",
"请总结这个时间序列：{ts_start} {ts_end}",
"请简要描述这个时间序列：{ts_start} {ts_end}",
"请简要概括这个时间序列：{ts_start} {ts_end}",
"请总结这个时间序列的主要特征：{ts_start} {ts_end}",
"请描述这个时间序列的主要表现：{ts_start} {ts_end}",
"请概括这个时间序列的整体情况：{ts_start} {ts_end}",
"请总结这个时间序列的整体特征：{ts_start} {ts_end}",
"请简述这个时间序列的主要模式：{ts_start} {ts_end}",
"请描述该时间序列：{ts_start} {ts_end}",
"请概括该时间序列：{ts_start} {ts_end}",
"请总结该时间序列：{ts_start} {ts_end}",
"请描述该时间序列的主要特征：{ts_start} {ts_end}",
"请概括该时间序列的整体特征：{ts_start} {ts_end}",
"请总结该时间序列的主要表现：{ts_start} {ts_end}",
"请对这个时间序列做简要描述：{ts_start} {ts_end}",
"请对这个时间序列做简要概括：{ts_start} {ts_end}",
"请对这个时间序列做简要总结：{ts_start} {ts_end}",
"请简要总结这个时间序列的主要情况：{ts_start} {ts_end}",
)
STAGE1_PROMPT_VARIANTS_BIVAR: tuple[str, ...] = (
"请描述这个双变量时间序列：{ts_start} {ts_end}",
"请概括这个双变量时间序列：{ts_start} {ts_end}",
"请总结这个双变量时间序列：{ts_start} {ts_end}",
"请简要描述这个双变量时间序列的整体表现：{ts_start} {ts_end}",
"请简要概括这个双变量时间序列：{ts_start} {ts_end}",
"请概括这个双变量时间序列的主要关系特征：{ts_start} {ts_end}",
"请总结这个双变量时间序列的整体变化与相互关系：{ts_start} {ts_end}",
"请总结这个双变量时间序列的主要特征：{ts_start} {ts_end}",
"请描述这个双变量时间序列的主要表现：{ts_start} {ts_end}",
"请概括这个双变量时间序列的整体情况：{ts_start} {ts_end}",
"请总结这个双变量时间序列的整体特征：{ts_start} {ts_end}",
"请简述这个双变量时间序列的主要关系模式：{ts_start} {ts_end}",
"请描述该双变量时间序列：{ts_start} {ts_end}",
"请概括该双变量时间序列：{ts_start} {ts_end}",
"请描述该双变量时间序列的主要表现：{ts_start} {ts_end}",
"请概括该双变量序列的整体特征：{ts_start} {ts_end}",
"请总结该双变量序列的主要模式：{ts_start} {ts_end}",
"请对这对时间序列做简要描述：{ts_start} {ts_end}",
"请对这对时间序列做简要概括：{ts_start} {ts_end}",
"请简要总结这对时间序列的主要情况：{ts_start} {ts_end}",
)
STAGE1_PROMPT_VARIANTS_MULTIVAR: tuple[str, ...] = (
"请描述这个多变量时间序列：{ts_start} {ts_end}",
"请概括这个多变量时间序列：{ts_start} {ts_end}",
"请总结这个多变量时间序列：{ts_start} {ts_end}",
"请简要描述这个多变量系统的整体表现：{ts_start} {ts_end}",
"请简要概括这个多变量时间序列：{ts_start} {ts_end}",
"请概括该多变量系统的主要结构特征：{ts_start} {ts_end}",
"请总结该多变量序列的整体变化模式：{ts_start} {ts_end}",
"请总结这个多变量时间序列的主要特征：{ts_start} {ts_end}",
"请描述这个多变量系统的主要表现：{ts_start} {ts_end}",
"请概括这个多变量系统的整体情况：{ts_start} {ts_end}",
"请总结这个多变量系统的整体特征：{ts_start} {ts_end}",
"请简述这个多变量系统的主要模式：{ts_start} {ts_end}",
"请描述该多变量时间序列：{ts_start} {ts_end}",
"请概括该多变量时间序列：{ts_start} {ts_end}",
"请描述该多变量时间序列的主要表现：{ts_start} {ts_end}",
"请概括该多变量时间序列的整体特征：{ts_start} {ts_end}",
"请总结该多变量系统的主要模式：{ts_start} {ts_end}",
"请对这个多变量时间序列做简要描述：{ts_start} {ts_end}",
"请对这个多变量时间序列做简要概括：{ts_start} {ts_end}",
"请简要总结这个多变量系统的主要情况：{ts_start} {ts_end}",
)
STAGE1_PROMPT_VARIANTS_BY_MODE: dict[PromptMode, tuple[str, ...]] = {
"univar": STAGE1_PROMPT_VARIANTS_UNIVAR,
"bivar": STAGE1_PROMPT_VARIANTS_BIVAR,
"multivar": STAGE1_PROMPT_VARIANTS_MULTIVAR,
}
STAGE2_SYSTEM_PROMPT = "你是专业的时间序列分析助手，请仅根据给定时间序列完成分析。"
STAGE2_PROMPT_FAMILIES_UNIVAR: dict[str, tuple[str, ...]] = {
"overall": (
"请对这个时间序列做深入分析：{ts_start} {ts_end}",
"请结合整体形态，对该时间序列做较完整的分析：{ts_start} {ts_end}",
"请围绕整体表现和变化脉络，对该序列进行深入分析：{ts_start} {ts_end}",
),
"pattern": (
"请分析该时间序列的主要模式及其相互关系：{ts_start} {ts_end}",
"请围绕趋势、周期性和局部波动，对该序列进行综合分析：{ts_start} {ts_end}",
"请从主要模式及其关联的角度，对该时间序列进行分析：{ts_start} {ts_end}",
),
"stability": (
"请分析这个时间序列的稳定性与可预测性：{ts_start} {ts_end}",
"请判断该序列的变化是否稳定，并说明其可预测性的来源：{ts_start} {ts_end}",
"请从稳定性和可预测性的角度，对该时间序列进行分析：{ts_start} {ts_end}",
),
"risk": (
"请分析该时间序列是否存在结构变化风险，并说明依据：{ts_start} {ts_end}",
"请从长期趋势、波动变化和潜在结构切换的角度分析该序列：{ts_start} {ts_end}",
"请评估该时间序列的结构风险，并结合整体变化给出分析：{ts_start} {ts_end}",
),
}
STAGE2_PROMPT_FAMILIES_BIVAR: dict[str, tuple[str, ...]] = {
"overall": (
"请对这个双变量时间序列做深入分析，重点概括整体变化与两条序列的关系：{ts_start} {ts_end}",
"请结合整体形态与变量间联系，对该双变量时间序列做较完整的分析：{ts_start} {ts_end}",
"请围绕整体表现、协同变化和差异关系，对这对时间序列进行深入分析：{ts_start} {ts_end}",
),
"pattern": (
"请分析该双变量时间序列的主要关系模式及其相互作用：{ts_start} {ts_end}",
"请围绕趋势一致性、节律同步和局部波动联动，对这对序列进行综合分析：{ts_start} {ts_end}",
"请从关系结构与模式特征的角度，对该双变量时间序列进行分析：{ts_start} {ts_end}",
),
"stability": (
"请分析这对时间序列关系结构的稳定性与可预测性：{ts_start} {ts_end}",
"请判断该双变量序列的协同变化是否稳定，并说明其可预测性的来源：{ts_start} {ts_end}",
"请从关系稳定性和联合可预测性的角度，对这对时间序列进行分析：{ts_start} {ts_end}",
),
"risk": (
"请分析该双变量时间序列是否存在关系结构变化风险，并说明依据：{ts_start} {ts_end}",
"请从长期趋势、波动联动和潜在结构切换的角度分析这对序列：{ts_start} {ts_end}",
"请评估这对时间序列的耦合风险，并结合整体变化给出分析：{ts_start} {ts_end}",
),
}
STAGE2_PROMPT_FAMILIES_MULTIVAR: dict[str, tuple[str, ...]] = {
"overall": (
"请对这个多变量时间序列系统做深入分析，重点概括整体结构与动态模式：{ts_start} {ts_end}",
"请结合系统整体形态与变量间协同关系，对该多变量时间序列做较完整的分析：{ts_start} {ts_end}",
"请围绕整体表现、系统结构和变量间互动，对该多变量序列进行深入分析：{ts_start} {ts_end}",
),
"pattern": (
"请分析该多变量时间序列系统的主要模式特征及其相互关系：{ts_start} {ts_end}",
"请围绕因子结构、同步协同、领先-滞后与局部波动，对该多变量系统进行综合分析：{ts_start} {ts_end}",
"请从系统模式与变量间关联的角度，对该多变量时间序列进行分析：{ts_start} {ts_end}",
),
"stability": (
"请分析这个多变量时间序列系统的结构稳定性与可预测性：{ts_start} {ts_end}",
"请判断该多变量系统的协同结构是否稳定，并说明其可预测性的来源：{ts_start} {ts_end}",
"请从系统稳定性和联合可预测性的角度，对该多变量时间序列进行分析：{ts_start} {ts_end}",
),
"risk": (
"请分析该多变量时间序列系统是否存在结构变化风险，并说明依据：{ts_start} {ts_end}",
"请从长期趋势、相关结构变化、波动联动和潜在状态切换的角度分析该多变量系统：{ts_start} {ts_end}",
"请评估该多变量时间序列的结构脆弱性与异常风险，并结合整体变化给出分析：{ts_start} {ts_end}",
),
}
STAGE2_PROMPT_FAMILIES_BY_MODE: dict[PromptMode, dict[str, tuple[str, ...]]] = {
"univar": STAGE2_PROMPT_FAMILIES_UNIVAR,
"bivar": STAGE2_PROMPT_FAMILIES_BIVAR,
"multivar": STAGE2_PROMPT_FAMILIES_MULTIVAR,
}
# Backward-compatible aliases: in the multivar package, the default exported prompt
# set should reflect the multivariate training path.
STAGE1_PROMPT_VARIANTS = STAGE1_PROMPT_VARIANTS_MULTIVAR
STAGE2_PROMPT_FAMILIES = STAGE2_PROMPT_FAMILIES_MULTIVAR
DEFAULT_STAGE2_PROMPT_FAMILY_WEIGHTS: dict[str, float] = {
"overall": 0.4,
"pattern": 0.25,
"stability": 0.2,
"risk": 0.15,
}
DEFAULT_STAGE2_LEVEL_WEIGHTS: dict[str, float] = {
# Phase 3 (2026-06-10): enable level_1/2 captions to give ScaleEncoder
# + LLM strong supervision for absolute (mu, sigma) decoding. Phase 2
# used only level_3/4 (analytical captions), starving the signal for
# simple-stats metrics (mean/std/min/max/median/...). See phase3_design
# §4.1.1. When the runner doesn't provide --level12-jsonl, train_stage1
# forces level_1/2 weights to 0 to preserve backward compatibility.
"level_1": 0.5,
"level_2": 0.3,
"level_3": 2.0,
"level_4": 1.0,
}
def register_ts_special_tokens(tokenizer, ts_token: str = "") -> int:
token_ids = register_alignment_special_tokens(
tokenizer,
ts_start_token=ts_token,
ts_end_token=TS_END_TOKEN,
scale_start_token=SCALE_START_TOKEN,
scale_end_token=SCALE_END_TOKEN,
)
return token_ids["ts_start"]
def register_alignment_special_tokens(
tokenizer,
*,
ts_start_token: str = TS_START_TOKEN,
ts_end_token: str = TS_END_TOKEN,
scale_start_token: str = SCALE_START_TOKEN,
scale_end_token: str = SCALE_END_TOKEN,
) -> dict[str, int]:
tokens = [ts_start_token, ts_end_token, scale_start_token, scale_end_token]
vocab = tokenizer.get_vocab()
missing_tokens = [token for token in tokens if token not in vocab]
if missing_tokens:
tokenizer.add_special_tokens({"additional_special_tokens": missing_tokens})
return {
"ts_start": tokenizer.convert_tokens_to_ids(ts_start_token),
"ts_end": tokenizer.convert_tokens_to_ids(ts_end_token),
"scale_start": tokenizer.convert_tokens_to_ids(scale_start_token),
"scale_end": tokenizer.convert_tokens_to_ids(scale_end_token),
}
def collate_fn(
batch: list[dict[str, Any]],
tokenizer,
ts_token: str = TS_START_TOKEN,
ts_end_token: str = TS_END_TOKEN,
ignore_index: int = -100,
max_length: int | None = None,
) -> dict[str, Any]:
if not batch:
raise ValueError("collate_fn requires a non-empty batch.")
special_token_ids = register_alignment_special_tokens(
tokenizer,
ts_start_token=ts_token,
ts_end_token=ts_end_token,
)
batch_type = _detect_batch_type(batch)
raw_ts_sequences: list[torch.Tensor] = []
channel_counts: list[int] = []
ts_lengths_list: list[int] = []
input_ids_list: list[torch.Tensor] = []
attention_masks: list[torch.Tensor] = []
labels_list: list[torch.Tensor] = []
text_valid_lengths: list[int] = []
ts_start_positions: list[int] = []
ts_end_positions: list[int] = []
for sample in batch:
raw_ts = _normalize_raw_ts(sample["raw_ts"])
raw_ts_sequences.append(raw_ts)
channel_counts.append(raw_ts.shape[0])
ts_lengths_list.append(raw_ts.shape[-1])
if batch_type == "text":
input_ids, attention_mask, labels = _build_text_training_example(
sample=sample,
tokenizer=tokenizer,
ignore_index=ignore_index,
max_length=max_length,
)
else:
input_ids, attention_mask, labels = _build_pretokenized_example(
sample=sample,
ignore_index=ignore_index,
max_length=max_length,
)
valid_length, start_pos, end_pos = _validate_single_placeholder(
input_ids=input_ids,
attention_mask=attention_mask,
ts_start_token_id=special_token_ids["ts_start"],
ts_end_token_id=special_token_ids["ts_end"],
)
text_valid_lengths.append(valid_length)
ts_start_positions.append(start_pos)
ts_end_positions.append(end_pos)
input_ids_list.append(input_ids)
attention_masks.append(attention_mask)
labels_list.append(labels)
batch_size = len(raw_ts_sequences)
max_channels = max(channel_counts)
max_ts_len = max(ts_lengths_list)
raw_ts = torch.zeros(
batch_size,
max_channels,
max_ts_len,
dtype=torch.float32,
)
raw_ts_channel_mask = torch.zeros(
batch_size,
max_channels,
dtype=torch.long,
)
raw_ts_attention_mask = torch.zeros(
batch_size,
max_ts_len,
dtype=torch.long,
)
for batch_index, sequence in enumerate(raw_ts_sequences):
n_channels, seq_len = sequence.shape
raw_ts[batch_index, :n_channels, :seq_len] = sequence
raw_ts_channel_mask[batch_index, :n_channels] = 1
raw_ts_attention_mask[batch_index, :seq_len] = 1
input_ids = pad_sequence(
input_ids_list,
batch_first=True,
padding_value=_get_pad_token_id(tokenizer),
)
attention_mask = pad_sequence(
attention_masks,
batch_first=True,
padding_value=0,
)
labels = pad_sequence(
labels_list,
batch_first=True,
padding_value=ignore_index,
)
return {
"raw_ts": raw_ts,
"raw_ts_channel_mask": raw_ts_channel_mask,
"raw_ts_attention_mask": raw_ts_attention_mask,
"text_valid_lengths": text_valid_lengths,
"ts_start_positions": ts_start_positions,
"ts_end_positions": ts_end_positions,
"input_ids": input_ids,
"attention_mask": attention_mask,
"labels": labels,
}
def get_stage1_prompt_variants(
ts_start_token: str = TS_START_TOKEN,
ts_end_token: str = TS_END_TOKEN,
prompt_variants: Sequence[str] | None = None,
prompt_mode: PromptMode = "multivar",
) -> list[str]:
variants = prompt_variants or STAGE1_PROMPT_VARIANTS_BY_MODE[prompt_mode]
formatted = [
_format_prompt_variant(
variant,
ts_start_token=ts_start_token,
ts_end_token=ts_end_token,
)
for variant in variants
]
if not formatted:
raise ValueError("At least one stage1 prompt variant is required.")
return formatted
def sample_stage1_prompt(
*,
ts_start_token: str = TS_START_TOKEN,
ts_end_token: str = TS_END_TOKEN,
prompt_variants: Sequence[str] | None = None,
prompt_mode: PromptMode = "multivar",
rng: random.Random | None = None,
) -> str:
variants = get_stage1_prompt_variants(
ts_start_token=ts_start_token,
ts_end_token=ts_end_token,
prompt_variants=prompt_variants,
prompt_mode=prompt_mode,
)
chooser = rng.choice if rng is not None else random.choice
return chooser(variants)
def build_stage1_training_samples(
record: dict[str, Any],
*,
ts_start_token: str = TS_START_TOKEN,
ts_end_token: str = TS_END_TOKEN,
prompt_variants: Sequence[str] | None = None,
rng: random.Random | None = None,
) -> list[dict[str, Any]]:
_validate_stage1_record(record)
samples: list[dict[str, Any]] = []
for level_key in ("level_1", "level_2"):
samples.append(
_build_stage1_sample(
record=record,
level_key=level_key,
ts_start_token=ts_start_token,
ts_end_token=ts_end_token,
prompt_variants=prompt_variants,
rng=rng,
)
)
return samples
class Stage1AlignmentDataset(torch.utils.data.Dataset):
def __init__(
self,
records: Sequence[dict[str, Any]],
*,
ts_start_token: str = TS_START_TOKEN,
ts_end_token: str = TS_END_TOKEN,
prompt_variants: Sequence[str] | None = None,
seed: int = 0,
dynamic: bool = True,
) -> None:
self.records = list(records)
self.ts_start_token = ts_start_token
self.ts_end_token = ts_end_token
self.prompt_variants = prompt_variants
self.seed = seed
self.dynamic = dynamic
self.epoch = 0
for record in self.records:
_validate_stage1_record(record)
def __len__(self) -> int:
return len(self.records) * 2
def set_epoch(self, epoch: int) -> None:
self.epoch = epoch
def __getitem__(self, index: int) -> dict[str, Any]:
if index < 0:
index += len(self)
if index < 0 or index >= len(self):
raise IndexError("Stage1AlignmentDataset index out of range.")
record_index, level_index = divmod(index, 2)
level_key = ("level_1", "level_2")[level_index]
sample_seed = self.seed + index
if self.dynamic:
sample_seed += self.epoch * max(len(self), 1)
rng = random.Random(sample_seed)
return _build_stage1_sample(
record=self.records[record_index],
level_key=level_key,
ts_start_token=self.ts_start_token,
ts_end_token=self.ts_end_token,
prompt_variants=self.prompt_variants,
rng=rng,
)
class Stage2AlignmentDataset(torch.utils.data.Dataset):
def __init__(
self,
records: Sequence[dict[str, Any]],
*,
system_prompt: str = STAGE2_SYSTEM_PROMPT,
ts_start_token: str = TS_START_TOKEN,
ts_end_token: str = TS_END_TOKEN,
prompt_families: dict[str, Sequence[str]] | None = None,
prompt_family_weights: dict[str, float] | None = None,
level_weights: dict[str, float] | None = None,
seed: int = 0,
dynamic: bool = True,
) -> None:
self.records = list(records)
self.system_prompt = system_prompt
self.ts_start_token = ts_start_token
self.ts_end_token = ts_end_token
self.prompt_families = prompt_families
self.prompt_family_weights = prompt_family_weights
self.level_weights = level_weights
self.seed = seed
self.dynamic = dynamic
self.epoch = 0
for record in self.records:
_validate_stage2_record(record)
def __len__(self) -> int:
return len(self.records)
def set_epoch(self, epoch: int) -> None:
self.epoch = epoch
def __getitem__(self, index: int) -> dict[str, Any]:
if index < 0:
index += len(self)
if index < 0 or index >= len(self):
raise IndexError("Stage2AlignmentDataset index out of range.")
sample_seed = self.seed + index
if self.dynamic:
sample_seed += self.epoch * max(len(self.records), 1)
rng = random.Random(sample_seed)
return build_stage2_training_sample(
self.records[index],
system_prompt=self.system_prompt,
ts_start_token=self.ts_start_token,
ts_end_token=self.ts_end_token,
prompt_families=self.prompt_families,
prompt_family_weights=self.prompt_family_weights,
level_weights=self.level_weights,
rng=rng,
)
class QAWarmupDataset(torch.utils.data.Dataset):
"""SFT warm-up dataset that consumes pre-built (prompt, answer) QA pairs.
Records are produced by build_qa_warmup_jsonl.py: each one already carries
a fully-rendered Chinese prompt (with ` ` placeholders) and the
deterministic reference answer (`metric_anchor = value;` joined by `；`).
"""
def __init__(
self,
records: Sequence[dict[str, Any]],
*,
system_prompt: str | None = None,
seed: int = 0,
dynamic: bool = False,
) -> None:
self.records = list(records)
self.system_prompt = system_prompt
self.seed = seed
self.dynamic = dynamic
self.epoch = 0
for record in self.records:
_validate_qa_warmup_record(record)
def __len__(self) -> int:
return len(self.records)
def set_epoch(self, epoch: int) -> None:
self.epoch = epoch
def __getitem__(self, index: int) -> dict[str, Any]:
if index < 0:
index += len(self)
if index < 0 or index >= len(self):
raise IndexError("QAWarmupDataset index out of range.")
record = self.records[index]
sample: dict[str, Any] = {
"raw_ts": record["raw_ts"],
"prompt": record["prompt"],
"target_text": record["answer"],
"target_level": "qa_warmup",
"source_id": record.get("id"),
}
# Prefer the record's own system_prompt (mixed pool: TSQA rows carry
# TSQA_SYSTEM_PROMPT); fall back to the dataset-level default for rows
# without one (our metric_qa) — identical to the previous behavior.
sp = record.get("system_prompt", self.system_prompt)
if sp is not None:
sample["system_prompt"] = sp
if "metric_keys" in record:
sample["metric_keys"] = record["metric_keys"]
return sample
def load_qa_warmup_records_from_jsonl(
*,
raw_values_path: str | Path,
qa_pairs_path: str | Path,
limit: int | None = None,
) -> list[dict[str, Any]]:
"""Join raw_values + qa_pairs by id. One record per QA pair (multiple per id)."""
raw_records = _read_jsonl_file(raw_values_path)
raw_by_id: dict[int, dict[str, Any]] = {}
for record in raw_records:
if "id" not in record:
continue
if not _has_valid_ts_values(record.get("values")):
continue
raw_by_id[int(record["id"])] = record
qa_records = _read_jsonl_file(qa_pairs_path, limit=limit)
records: list[dict[str, Any]] = []
for qa in qa_records:
record_id = qa.get("id")
if record_id is None:
continue
raw = raw_by_id.get(int(record_id))
if raw is None:
continue
prompt = qa.get("prompt")
answer = qa.get("answer")
if not isinstance(prompt, str) or not isinstance(answer, str):
continue
if not prompt.strip() or not answer.strip():
continue
rec = {
"id": int(record_id),
"qa_index": qa.get("qa_index"),
"raw_ts": raw["values"],
"prompt": prompt,
"answer": answer,
"metric_keys": qa.get("metric_keys", []),
}
# Carry per-row system_prompt / task_type through so a mixed warmup pool
# (metric_qa + caption + TSQA) keeps each task's own system prompt. Rows
# without these fields (our metric_qa sft_pairs) fall back to the
# dataset-level default in QAWarmupDataset — behavior unchanged.
if qa.get("system_prompt") is not None:
rec["system_prompt"] = qa["system_prompt"]
if qa.get("task_type") is not None:
rec["task_type"] = qa["task_type"]
records.append(rec)
return records
def _validate_qa_warmup_record(record: dict[str, Any]) -> None:
missing = [key for key in ("raw_ts", "prompt", "answer") if key not in record]
if missing:
raise ValueError(
f"QAWarmup record is missing required fields: {', '.join(missing)}."
)
if not _has_valid_ts_values(record["raw_ts"]):
raise ValueError("QAWarmup record raw_ts must be a non-empty sequence without null values.")
prompt = record["prompt"]
if not isinstance(prompt, str) or TS_START_TOKEN not in prompt or TS_END_TOKEN not in prompt:
raise ValueError(
f"QAWarmup prompt must contain both {TS_START_TOKEN} and {TS_END_TOKEN}."
)
def build_univariate_stage1_records(
sample_records: Sequence[dict[str, Any]],
level_records: Sequence[dict[str, Any]],
) -> list[dict[str, Any]]:
samples_by_id = {
int(record["id"]): record
for record in sample_records
}
levels_by_id = OrderedDict(
(int(record["id"]), record)
for record in level_records
)
records: list[dict[str, Any]] = []
for record_id, level_record in levels_by_id.items():
if record_id not in samples_by_id:
raise ValueError(f"Missing sample values for record id {record_id}.")
sample_record = samples_by_id[record_id]
if not _has_valid_ts_values(sample_record.get("values")):
continue
records.append(
{
"id": record_id,
"dataset": sample_record.get("dataset"),
"channel": sample_record.get("channel"),
"raw_ts": sample_record["values"],
"level_1": (
level_record.get("level_1_revised")
or level_record.get("original_level_1")
or level_record.get("level_1")
),
"level_2": (
level_record.get("level_2_revised")
or level_record.get("original_level_2")
or level_record.get("level_2")
),
}
)
return records
def build_univariate_stage2_records(
sample_records: Sequence[dict[str, Any]],
level3_records: Sequence[dict[str, Any]],
level4_records: Sequence[dict[str, Any]],
level12_records: Sequence[dict[str, Any]] | None = None,
) -> list[dict[str, Any]]:
samples_by_id = {
int(record["id"]): record
for record in sample_records
}
level3_by_id = OrderedDict(
(int(record["id"]), record)
for record in level3_records
)
level4_by_id = {
int(record["id"]): record
for record in level4_records
}
# Phase 3: optional level_1/level_2 supply. Records only get level_1/level_2
# populated when level12_records is provided; otherwise the sampler in
# train_stage1 forces their weights to 0 so the missing field is never read.
level12_by_id: dict[int, dict[str, Any]] = {}
if level12_records is not None:
level12_by_id = {
int(record["id"]): record
for record in level12_records
}
records: list[dict[str, Any]] = []
for record_id, level3_record in level3_by_id.items():
if record_id not in samples_by_id:
raise ValueError(f"Missing sample values for record id {record_id}.")
if record_id not in level4_by_id:
raise ValueError(f"Missing level_4 text for record id {record_id}.")
sample_record = samples_by_id[record_id]
if not _has_valid_ts_values(sample_record.get("values")):
continue
level4_record = level4_by_id[record_id]
record: dict[str, Any] = {
"id": record_id,
"dataset": sample_record.get("dataset"),
"channel": sample_record.get("channel"),
"raw_ts": sample_record["values"],
"level_3": level3_record["level_3"],
"level_4": level4_record["level_4"],
"level_3_prompt": level3_record.get("prompt"),
"level_4_prompt": level4_record.get("prompt"),
}
level12_record = level12_by_id.get(record_id)
if level12_record is not None:
lvl1 = (
level12_record.get("level_1_revised")
or level12_record.get("original_level_1")
or level12_record.get("level_1")
)
lvl2 = (
level12_record.get("level_2_revised")
or level12_record.get("original_level_2")
or level12_record.get("level_2")
)
if lvl1 is not None:
record["level_1"] = lvl1
record["level_1_prompt"] = level12_record.get("level_1_prompt") or level12_record.get("prompt")
if lvl2 is not None:
record["level_2"] = lvl2
record["level_2_prompt"] = level12_record.get("level_2_prompt") or level12_record.get("prompt")
records.append(record)
return records
def load_stage1_records_from_jsonl(
*,
samples_path: str | Path,
level12_path: str | Path,
limit: int | None = None,
) -> list[dict[str, Any]]:
sample_records = _read_jsonl_file(samples_path, limit=limit)
level_records = _read_jsonl_file(level12_path, limit=limit)
return build_univariate_stage1_records(sample_records, level_records)
def load_stage1_records_from_univar_tar(
archive_path: str | Path,
*,
limit: int | None = None,
) -> list[dict[str, Any]]:
sample_records = _read_jsonl_from_tar_zst(
archive_path,
member_path="univar/samples.jsonl",
limit=limit,
)
level_records = _read_jsonl_from_tar_zst(
archive_path,
member_path="univar/level12.jsonl",
limit=limit,
)
return build_univariate_stage1_records(sample_records, level_records)
def load_stage2_records_from_jsonl(
*,
samples_path: str | Path,
level3_path: str | Path,
level4_path: str | Path,
level12_path: str | Path | None = None,
limit: int | None = None,
) -> list[dict[str, Any]]:
sample_records = _read_jsonl_file(samples_path, limit=limit)
level3_records = _read_jsonl_file(level3_path, limit=limit)
level4_records = _read_jsonl_file(level4_path, limit=limit)
level12_records = (
_read_jsonl_file(level12_path, limit=limit) if level12_path else None
)
return build_univariate_stage2_records(
sample_records,
level3_records,
level4_records,
level12_records=level12_records,
)
def load_stage2_records_from_univar_tar(
samples_archive_path: str | Path,
level34_archive_path: str | Path,
*,
limit: int | None = None,
) -> list[dict[str, Any]]:
sample_records = _read_jsonl_from_tar_zst(
samples_archive_path,
member_path="univar/samples.jsonl",
limit=limit,
)
level3_records = _read_jsonl_from_tar_zst(
level34_archive_path,
member_path="univar_level_34/level3.jsonl",
limit=limit,
)
level4_records = _read_jsonl_from_tar_zst(
level34_archive_path,
member_path="univar_level_34/level4.jsonl",
limit=limit,
)
return build_univariate_stage2_records(sample_records, level3_records, level4_records)
def _detect_batch_type(batch: list[dict[str, Any]]) -> str:
has_text = [("prompt" in sample and "target_text" in sample) for sample in batch]
has_tokens = [("input_ids" in sample and "labels" in sample) for sample in batch]
if all(has_text) and not any(has_tokens):
return "text"
if all(has_tokens) and not any(has_text):
return "tokenized"
raise ValueError(
"Batch must contain either only prompt/target_text samples or only pretokenized samples."
)
def _build_stage1_sample(
*,
record: dict[str, Any],
level_key: str,
ts_start_token: str,
ts_end_token: str,
prompt_variants: Sequence[str] | None,
rng: random.Random | None,
) -> dict[str, Any]:
prompt_mode = _infer_prompt_mode(record["raw_ts"])
return {
"raw_ts": record["raw_ts"],
"prompt": sample_stage1_prompt(
ts_start_token=ts_start_token,
ts_end_token=ts_end_token,
prompt_variants=prompt_variants,
prompt_mode=prompt_mode,
rng=rng,
),
"target_text": record[level_key],
"target_level": level_key,
"source_id": record.get("id"),
}
def get_stage2_prompt_families(
ts_start_token: str = TS_START_TOKEN,
ts_end_token: str = TS_END_TOKEN,
prompt_families: dict[str, Sequence[str]] | None = None,
prompt_mode: PromptMode = "multivar",
) -> dict[str, list[str]]:
families = prompt_families or STAGE2_PROMPT_FAMILIES_BY_MODE[prompt_mode]
formatted = {
family: [
_format_prompt_variant(
variant,
ts_start_token=ts_start_token,
ts_end_token=ts_end_token,
)
for variant in variants
]
for family, variants in families.items()
}
if not formatted:
raise ValueError("At least one stage2 prompt family is required.")
for family, variants in formatted.items():
if not variants:
raise ValueError(f"Stage2 prompt family '{family}' must contain at least one prompt.")
return formatted
def sample_stage2_prompt(
*,
ts_start_token: str = TS_START_TOKEN,
ts_end_token: str = TS_END_TOKEN,
prompt_families: dict[str, Sequence[str]] | None = None,
prompt_family_weights: dict[str, float] | None = None,
prompt_mode: PromptMode = "multivar",
rng: random.Random | None = None,
) -> tuple[str, str]:
formatted_families = get_stage2_prompt_families(
ts_start_token=ts_start_token,
ts_end_token=ts_end_token,
prompt_families=prompt_families,
prompt_mode=prompt_mode,
)
family_weights = _resolve_weight_mapping(
prompt_family_weights,
defaults=DEFAULT_STAGE2_PROMPT_FAMILY_WEIGHTS,
allowed_keys=formatted_families.keys(),
mapping_name="stage2 prompt family weights",
)
chooser = rng or random
family = _weighted_choice(family_weights, chooser)
return family, chooser.choice(formatted_families[family])
def sample_stage2_target_level(
*,
level_weights: dict[str, float] | None = None,
rng: random.Random | None = None,
) -> str:
chooser = rng or random
weights = _resolve_weight_mapping(
level_weights,
defaults=DEFAULT_STAGE2_LEVEL_WEIGHTS,
allowed_keys=DEFAULT_STAGE2_LEVEL_WEIGHTS.keys(),
mapping_name="stage2 target weights",
)
return _weighted_choice(weights, chooser)
def build_stage2_training_sample(
record: dict[str, Any],
*,
system_prompt: str = STAGE2_SYSTEM_PROMPT,
ts_start_token: str = TS_START_TOKEN,
ts_end_token: str = TS_END_TOKEN,
prompt_families: dict[str, Sequence[str]] | None = None,
prompt_family_weights: dict[str, float] | None = None,
level_weights: dict[str, float] | None = None,
rng: random.Random | None = None,
) -> dict[str, Any]:
_validate_stage2_record(record)
prompt_mode = _infer_prompt_mode(record["raw_ts"])
target_level = sample_stage2_target_level(level_weights=level_weights, rng=rng)
if target_level in ("level_1", "level_2"):
# Phase 3: level_1/2 are short basic-stat captions. 8cd82fa enabled the
# data fields + weights but never wired a prompt path for them, so the
# aligned-family machinery (DIRECT_STAGE2_SOURCE_PROMPT_RULES /
# ALIGNED_STAGE2_PROMPT_VARIANTS) only covers level_3/4 → KeyError when
# level_1/2 is sampled. Reuse the Stage 1 generic describe prompts
# (STAGE1_PROMPT_VARIANTS_BY_MODE), which match the short-caption task.
source_prompt = None
family = "describe"
prompt = sample_stage1_prompt(
ts_start_token=ts_start_token,
ts_end_token=ts_end_token,
prompt_mode=prompt_mode,
rng=rng,
)
elif (source_prompt := _get_stage2_source_prompt(record, target_level)):
family, prompt = build_aligned_stage2_prompt(
source_prompt=source_prompt,
target_level=target_level,
ts_start_token=ts_start_token,
ts_end_token=ts_end_token,
prompt_families=prompt_families,
prompt_mode=prompt_mode,
rng=rng,
)
else:
family, prompt = sample_stage2_prompt(
ts_start_token=ts_start_token,
ts_end_token=ts_end_token,
prompt_families=prompt_families,
prompt_family_weights=prompt_family_weights,
prompt_mode=prompt_mode,
rng=rng,
)
# Phase 3 (2026-06-10): when env TS_ALIGN_PROMPT_TASK_TYPE=1, prepend a
# "caption" prefix. Pairs with the QA-side prefix
# injected by build_single_metric_prompt so the model can disambiguate
# which reward branch it's being trained on. Env-gated for backward
# compat; Phase 2 SFT builds remain unchanged.
from .qa_templates import task_type_token_prefix
task_type_prefix = task_type_token_prefix("caption")
if task_type_prefix:
prompt = task_type_prefix + prompt
return {
"raw_ts": record["raw_ts"],
"system_prompt": system_prompt,
"prompt": prompt,
"target_text": record[target_level],
"target_level": target_level,
"prompt_family": family,
"source_prompt": source_prompt,
"source_id": record.get("id"),
}
def _validate_stage1_record(record: dict[str, Any]) -> None:
missing = [key for key in ("raw_ts", "level_1", "level_2") if key not in record]
if missing:
raise ValueError(
f"Stage1 record is missing required fields: {', '.join(missing)}."
)
if not _has_valid_ts_values(record["raw_ts"]):
raise ValueError("Stage1 record raw_ts must be a non-empty sequence without null values.")
def _validate_stage2_record(record: dict[str, Any]) -> None:
missing = [key for key in ("raw_ts", "level_3", "level_4") if key not in record]
if missing:
raise ValueError(
f"Stage2 record is missing required fields: {', '.join(missing)}."
)
if not _has_valid_ts_values(record["raw_ts"]):
raise ValueError("Stage2 record raw_ts must be a non-empty sequence without null values.")
def _has_valid_ts_values(values: Any) -> bool:
if values is None:
return False
try:
tensor = _normalize_raw_ts(values)
except (TypeError, ValueError):
return False
return tensor.numel() > 0 and torch.isfinite(tensor).all().item()
def _infer_prompt_mode(raw_ts: Any) -> PromptMode:
tensor = _normalize_raw_ts(raw_ts)
n_channels = int(tensor.shape[0])
if n_channels <= 1:
return "univar"
if n_channels == 2:
return "bivar"
return "multivar"
def _get_stage2_source_prompt(record: dict[str, Any], target_level: str) -> str | None:
prompt_key = f"{target_level}_prompt"
source_prompt = record.get(prompt_key)
if isinstance(source_prompt, str) and source_prompt.strip():
return source_prompt.strip()
return None
def build_aligned_stage2_prompt(
*,
source_prompt: str,
target_level: str,
ts_start_token: str = TS_START_TOKEN,
ts_end_token: str = TS_END_TOKEN,
prompt_families: dict[str, Sequence[str]] | None = None,
prompt_mode: PromptMode = "multivar",
rng: random.Random | None = None,
) -> tuple[str, str]:
source_topic = infer_stage2_source_topic(
source_prompt=source_prompt,
target_level=target_level,
prompt_mode=prompt_mode,
)
family = SOURCE_TOPIC_TO_PROMPT_FAMILY[source_topic]
variants = get_aligned_stage2_prompt_variants(
prompt_mode=prompt_mode,
source_topic=source_topic,
ts_start_token=ts_start_token,
ts_end_token=ts_end_token,
)
chooser = rng.choice if rng is not None else random.choice
return family, chooser(variants)
def infer_stage2_prompt_family_from_source_prompt(
*,
source_prompt: str,
target_level: str,
prompt_mode: PromptMode = "multivar",
) -> str:
source_topic = infer_stage2_source_topic(
source_prompt=source_prompt,
target_level=target_level,
prompt_mode=prompt_mode,
)
return SOURCE_TOPIC_TO_PROMPT_FAMILY[source_topic]
def infer_stage2_source_topic(
*,
source_prompt: str,
target_level: str,
prompt_mode: PromptMode = "multivar",
) -> str:
headline = _extract_source_prompt_headline(source_prompt)
rules = DIRECT_STAGE2_SOURCE_PROMPT_RULES_BY_MODE[prompt_mode][target_level]
for phrase, source_topic in rules:
if phrase in headline:
return source_topic
raise ValueError(
f"Unable to map source prompt headline to aligned topic for mode={prompt_mode}, "
f"target_level={target_level}, headline={headline!r}."
)
def get_aligned_stage2_prompt_variants(
*,
prompt_mode: PromptMode,
source_topic: str,
ts_start_token: str = TS_START_TOKEN,
ts_end_token: str = TS_END_TOKEN,
) -> list[str]:
prompt_variants = ALIGNED_STAGE2_PROMPT_VARIANTS_BY_MODE[prompt_mode][source_topic]
return [
_format_prompt_variant(
variant,
ts_start_token=ts_start_token,
ts_end_token=ts_end_token,
)
for variant in prompt_variants
]
def _contains_any(text: str, keywords: Sequence[str]) -> bool:
return any(keyword in text for keyword in keywords)
def _extract_source_prompt_headline(source_prompt: str) -> str:
for line in source_prompt.splitlines():
headline = line.strip()
if headline:
return headline
return source_prompt.strip()
SOURCE_TOPIC_TO_PROMPT_FAMILY: dict[str, str] = {
"pattern": "pattern",
"association": "pattern",
"summary": "overall",
"coupling_stability": "stability",
"frequency": "stability",
"risk": "risk",
"deep_summary": "overall",
}
DIRECT_STAGE2_SOURCE_PROMPT_RULES_UNIVAR: dict[str, tuple[tuple[str, str], ...]] = {
"level_3": (
("综合分析其主要模式", "pattern"),
("主要模式特征", "pattern"),
("模式总结", "summary"),
("关联关系", "association"),
("核心模式", "summary"),
),
"level_4": (
("稳定性、复杂度和结构风险", "coupling_stability"),
("稳定性和可预测性", "coupling_stability"),
("频率结构与复杂度", "frequency"),
("异常风险与结构脆弱性", "risk"),
("最具技术价值的深层特征", "deep_summary"),
),
}
DIRECT_STAGE2_SOURCE_PROMPT_RULES_BIVAR: dict[str, tuple[tuple[str, str], ...]] = {
"level_3": (
("综合分析两序列的关系模式", "pattern"),
("主要关系模式", "pattern"),
("关系总结", "summary"),
("内在关联", "association"),
("最突出的关系特征", "summary"),
),
"level_4": (
("综合分析两序列的深层耦合特性", "deep_summary"),
("耦合结构与时变稳定性", "coupling_stability"),
("因果结构与频域特征", "frequency"),
("结构脆弱性与极端联动风险", "risk"),
("最具技术价值的深层特征", "deep_summary"),
),
}
DIRECT_STAGE2_SOURCE_PROMPT_RULES_MULTIVAR: dict[str, tuple[tuple[str, str], ...]] = {
"level_3": (
("综合分析其结构和动态模式", "pattern"),
("主要模式特征", "pattern"),
("总结概括", "summary"),
("关联关系", "association"),
("核心特征", "summary"),
),
"level_4": (
("综合分析系统的动态耦合特性和结构稳定性", "deep_summary"),
("动态耦合结构与时变稳定性", "coupling_stability"),
("频率特征与季节性结构", "frequency"),
("结构脆弱性与异常特征", "risk"),
("最具技术价值的深层特征", "deep_summary"),
),
}
DIRECT_STAGE2_SOURCE_PROMPT_RULES_BY_MODE: dict[
PromptMode, dict[str, tuple[tuple[str, str], ...]]
] = {
"univar": DIRECT_STAGE2_SOURCE_PROMPT_RULES_UNIVAR,
"bivar": DIRECT_STAGE2_SOURCE_PROMPT_RULES_BIVAR,
"multivar": DIRECT_STAGE2_SOURCE_PROMPT_RULES_MULTIVAR,
}
ALIGNED_STAGE2_PROMPT_VARIANTS_UNIVAR: dict[str, tuple[str, ...]] = {
"pattern": (
"请综合分析这个时间序列的主要模式：{ts_start} {ts_end}",
"请识别这个时间序列的主要模式特征：{ts_start} {ts_end}",
"请分析这个时间序列的主要模式：{ts_start} {ts_end}",
),
"association": (
"请分析这个时间序列各特征之间的关联关系：{ts_start} {ts_end}",
"请分析这个时间序列的关联关系：{ts_start} {ts_end}",
),
"summary": (
"请对这个时间序列做模式总结：{ts_start} {ts_end}",
"请概括这个时间序列的核心模式：{ts_start} {ts_end}",
"请总结这个时间序列的模式：{ts_start} {ts_end}",
),
"coupling_stability": (
"请分析这个时间序列的稳定性、复杂度和结构风险：{ts_start} {ts_end}",
"请重点分析这个时间序列的稳定性和可预测性：{ts_start} {ts_end}",
"请分析这个时间序列的稳定性和可预测性：{ts_start} {ts_end}",
),
"frequency": (
"请深入分析这个时间序列的频率结构与复杂度：{ts_start} {ts_end}",
"请分析这个时间序列的频率结构与复杂度：{ts_start} {ts_end}",
),
"risk": (
"请评估这个时间序列的异常风险与结构脆弱性：{ts_start} {ts_end}",
"请评估这个时间序列的结构脆弱性与异常风险：{ts_start} {ts_end}",
),
"deep_summary": (
"请概括这个时间序列最具技术价值的深层特征：{ts_start} {ts_end}",
"请概括这个时间序列的深层特征：{ts_start} {ts_end}",
),
}
ALIGNED_STAGE2_PROMPT_VARIANTS_BIVAR: dict[str, tuple[str, ...]] = {
"pattern": (
"请综合分析这对时间序列的关系模式：{ts_start} {ts_end}",
"请识别这对时间序列的主要关系模式：{ts_start} {ts_end}",
"请分析这对时间序列的关系模式：{ts_start} {ts_end}",
),
"association": (
"请分析这对时间序列各维度特征之间的内在关联：{ts_start} {ts_end}",
"请分析这对时间序列的内在关联：{ts_start} {ts_end}",
),
"summary": (
"请对这对时间序列做关系总结：{ts_start} {ts_end}",
"请概括这对时间序列最突出的关系特征：{ts_start} {ts_end}",
"请总结这对时间序列的关系特征：{ts_start} {ts_end}",
),
"coupling_stability": (
"请分析该双变量时间序列的耦合结构与时变稳定性：{ts_start} {ts_end}",
"请重点分析这对时间序列的耦合结构与时变稳定性：{ts_start} {ts_end}",
"请分析这对时间序列的耦合结构与时变稳定性：{ts_start} {ts_end}",
),
"frequency": (
"请分析该双变量时间序列的因果结构与频域特征：{ts_start} {ts_end}",
"请深入分析该双变量时间序列的因果结构与频域特征：{ts_start} {ts_end}",
"请分析这对时间序列的因果结构与频域特征：{ts_start} {ts_end}",
),
"risk": (
"请评估该双变量时间序列的结构脆弱性与极端联动风险：{ts_start} {ts_end}",
"请评估这对时间序列的结构脆弱性与极端联动风险：{ts_start} {ts_end}",
),
"deep_summary": (
"请概括该双变量时间序列最具技术价值的深层特征：{ts_start} {ts_end}",
"请概括这对时间序列最具技术价值的深层特征：{ts_start} {ts_end}",
),
}
ALIGNED_STAGE2_PROMPT_VARIANTS_MULTIVAR: dict[str, tuple[str, ...]] = {
"pattern": (
"请综合分析该多变量时间序列的结构和动态模式：{ts_start} {ts_end}",
"请识别该多变量系统的主要模式特征：{ts_start} {ts_end}",
"请分析该多变量时间序列的结构和动态模式：{ts_start} {ts_end}",
),
"association": (
"请分析该多变量系统各特征之间的关联关系：{ts_start} {ts_end}",
"请分析该多变量时间序列的关联关系：{ts_start} {ts_end}",
),
"summary": (
"请对该多变量系统做总结概括：{ts_start} {ts_end}",
"请概括该多变量系统的核心特征：{ts_start} {ts_end}",
"请总结该多变量时间序列的特征：{ts_start} {ts_end}",
),
"coupling_stability": (
"请分析该多变量时间序列的动态耦合特性和结构稳定性：{ts_start} {ts_end}",
"请分析该多变量系统的动态耦合结构与时变稳定性：{ts_start} {ts_end}",
"请重点分析该多变量系统的动态耦合结构与时变稳定性：{ts_start} {ts_end}",
"请分析该多变量系统的动态耦合与结构稳定性：{ts_start} {ts_end}",
),
"frequency": (
"请分析该多变量时间序列的频率特征与季节性结构：{ts_start} {ts_end}",
"请深入分析该多变量时间序列的频率特征与季节性结构：{ts_start} {ts_end}",
"请分析该多变量系统的频率特征与季节性结构：{ts_start} {ts_end}",
),
"risk": (
"请评估该多变量时间序列的结构脆弱性与异常特征：{ts_start} {ts_end}",
"请评估该多变量系统的结构脆弱性与异常特征：{ts_start} {ts_end}",
),
"deep_summary": (
"请概括该多变量时间序列最具技术价值的深层特征：{ts_start} {ts_end}",
"请概括该多变量系统最具技术价值的深层特征：{ts_start} {ts_end}",
),
}
ALIGNED_STAGE2_PROMPT_VARIANTS_BY_MODE: dict[PromptMode, dict[str, tuple[str, ...]]] = {
"univar": ALIGNED_STAGE2_PROMPT_VARIANTS_UNIVAR,
"bivar": ALIGNED_STAGE2_PROMPT_VARIANTS_BIVAR,
"multivar": ALIGNED_STAGE2_PROMPT_VARIANTS_MULTIVAR,
}
def _normalize_raw_ts(raw_ts: Any) -> torch.Tensor:
tensor = torch.as_tensor(raw_ts, dtype=torch.float32)
if tensor.ndim == 1:
return tensor.unsqueeze(0)
if tensor.ndim == 2:
if tensor.shape[0] == 1:
return tensor
if tensor.shape[1] == 1:
return tensor.transpose(0, 1)
# For multivariate inputs, prefer [C, L]. If the first dimension is much
# larger, interpret the tensor as [L, C] and transpose into channel-first layout.
if tensor.shape[0] > tensor.shape[1]:
return tensor.transpose(0, 1)
return tensor
raise ValueError("raw_ts must have shape [L], [C, L], or [L, C].")
def _build_text_training_example(
*,
sample: dict[str, Any],
tokenizer,
ignore_index: int,
max_length: int | None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
prompt_messages: list[dict[str, str]] = []
system_prompt = sample.get("system_prompt")
if system_prompt is not None:
prompt_messages.append({"role": "system", "content": system_prompt})
prompt_messages.append({"role": "user", "content": sample["prompt"]})
full_messages = [
*prompt_messages,
{"role": "assistant", "content": sample["target_text"]},
]
prompt_ids = _apply_chat_template(
tokenizer,
prompt_messages,
add_generation_prompt=True,
)
input_ids = _apply_chat_template(
tokenizer,
full_messages,
add_generation_prompt=False,
)
if len(prompt_ids) >= len(input_ids):
raise ValueError("Chat template must leave assistant tokens after the user prompt prefix.")
attention_mask = [1] * len(input_ids)
labels = [ignore_index] * len(prompt_ids) + input_ids[len(prompt_ids) :]
input_ids_tensor = torch.tensor(input_ids, dtype=torch.long)
attention_mask_tensor = torch.tensor(attention_mask, dtype=torch.long)
labels_tensor = torch.tensor(labels, dtype=torch.long)
return _truncate_text_fields(
input_ids=input_ids_tensor,
attention_mask=attention_mask_tensor,
labels=labels_tensor,
max_length=max_length,
)
def _build_pretokenized_example(
*,
sample: dict[str, Any],
ignore_index: int,
max_length: int | None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
input_ids = torch.as_tensor(sample["input_ids"], dtype=torch.long)
attention_mask = torch.as_tensor(
sample.get("attention_mask", torch.ones_like(input_ids)),
dtype=torch.long,
)
labels = torch.as_tensor(sample["labels"], dtype=torch.long)
if input_ids.ndim != 1 or attention_mask.ndim != 1 or labels.ndim != 1:
raise ValueError("Pretokenized input_ids, attention_mask, and labels must be 1D.")
return _truncate_text_fields(
input_ids=input_ids,
attention_mask=attention_mask,
labels=labels,
max_length=max_length,
)
def _truncate_text_fields(
*,
input_ids: torch.Tensor,
attention_mask: torch.Tensor,
labels: torch.Tensor,
max_length: int | None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
if max_length is None or input_ids.numel() <= max_length:
return input_ids, attention_mask, labels
input_ids = input_ids[:max_length]
attention_mask = attention_mask[:max_length]
labels = labels[:max_length]
return input_ids, attention_mask, labels
def _apply_chat_template(
tokenizer,
messages: list[dict[str, str]],
*,
add_generation_prompt: bool,
) -> list[int]:
if not hasattr(tokenizer, "apply_chat_template"):
raise ValueError("Tokenizer must support apply_chat_template for stage1 text examples.")
input_ids = tokenizer.apply_chat_template(
messages,
tokenize=True,
add_generation_prompt=add_generation_prompt,
)
if isinstance(input_ids, torch.Tensor):
return input_ids.tolist()
return list(input_ids)
def _validate_single_placeholder(
*,
input_ids: torch.Tensor,
attention_mask: torch.Tensor,
ts_start_token_id: int,
ts_end_token_id: int,
) -> tuple[int, int, int]:
valid_length = int(attention_mask.sum().item())
valid_input_ids = input_ids[:valid_length]
start_positions = (valid_input_ids == ts_start_token_id).nonzero(as_tuple=False).flatten()
end_positions = (valid_input_ids == ts_end_token_id).nonzero(as_tuple=False).flatten()
if start_positions.numel() != 1 or end_positions.numel() != 1:
raise ValueError("Each sample must contain exactly one and one token.")
start_pos = int(start_positions.item())
end_pos = int(end_positions.item())
if start_pos >= end_pos:
raise ValueError(" must appear before in each sample.")
return valid_length, start_pos, end_pos
def _format_prompt_variant(
variant: str,
*,
ts_start_token: str,
ts_end_token: str,
) -> str:
formatted = variant.format(
ts_start=ts_start_token,
ts_end=ts_end_token,
)
if formatted.count(ts_start_token) != 1 or formatted.count(ts_end_token) != 1:
raise ValueError(
"Each stage1 prompt variant must contain exactly one and one placeholder."
)
return formatted
def _resolve_weight_mapping(
weights: dict[str, float] | None,
*,
defaults: dict[str, float],
allowed_keys,
mapping_name: str,
) -> dict[str, float]:
resolved = dict(defaults)
if weights is not None:
resolved.update(weights)
resolved = {key: float(value) for key, value in resolved.items() if key in allowed_keys}
if not resolved:
raise ValueError(f"{mapping_name} must contain at least one entry.")
if any(value < 0 for value in resolved.values()):
raise ValueError(f"{mapping_name} cannot contain negative weights.")
total = sum(resolved.values())
if total <= 0:
raise ValueError(f"{mapping_name} must sum to a positive value.")
return resolved
def _weighted_choice(weights: dict[str, float], rng: random.Random) -> str:
total = sum(weights.values())
threshold = rng.random() * total
cumulative = 0.0
last_key = next(iter(weights))
for key, value in weights.items():
cumulative += value
last_key = key
if threshold <= cumulative:
return key
return last_key
def _read_jsonl_file(path: str | Path, *, limit: int | None = None) -> list[dict[str, Any]]:
import json
records: list[dict[str, Any]] = []
with open(path, "r", encoding="utf-8") as handle:
for index, line in enumerate(handle):
if limit is not None and index >= limit:
break
records.append(json.loads(line))
return records
def _read_jsonl_from_tar_zst(
archive_path: str | Path,
*,
member_path: str,
limit: int | None = None,
) -> list[dict[str, Any]]:
import json
command = (
f"zstd -dc {Path(archive_path)} | tar -xOf - {member_path}"
)
process = subprocess.Popen(
["bash", "-lc", command],
text=True,
encoding="utf-8",
errors="replace",
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
bufsize=1,
)
records: list[dict[str, Any]] = []
assert process.stdout is not None
assert process.stderr is not None
try:
for index, line in enumerate(process.stdout):
if limit is not None and index >= limit:
process.terminate()
break
if not line.strip():
continue
records.append(json.loads(line))
finally:
process.stdout.close()
stderr = process.stderr.read()
process.wait()
if process.returncode not in (0, -15):
raise RuntimeError(f"Failed to extract {member_path} from {archive_path}: {stderr}")
return records
def _get_pad_token_id(tokenizer) -> int:
pad_token_id = getattr(tokenizer, "pad_token_id", None)
if pad_token_id is not None:
return int(pad_token_id)
eos_token_id = getattr(tokenizer, "eos_token_id", None)
if eos_token_id is not None:
return int(eos_token_id)
raise ValueError("Tokenizer must define pad_token_id or eos_token_id for collation.")
# ---------------------------------------------------------------------------
# Phase 3 Tier-3 §T0 (2026-06-10): resumable samplers with persistent
# (epoch, position) state, so --resume-from sees the same prompt sequence
# the pre-kill trajectory would have. Closes the only remaining post-resume
# divergence source after Tier-2 §5/§6 (trainer_state + RNG persistence) —
# without T0, even with model weights + optimizer + RNG perfectly restored
# the DataLoader would yield the FIRST batch of a fresh epoch instead of
# resuming mid-epoch, and the next ~50 micro-batches would consume entirely
# different prompts than the original run.
#
# Both classes save {epoch: int, position_in_epoch: int}; load_state_dict
# resumes from the saved position; set_epoch (called explicitly at epoch
# boundary by the train loop) resets the position. Single-rank and
# distributed variants share the same state-dict shape so train_stage3
# can save without caring which is in use.
# ---------------------------------------------------------------------------
class _ResumableSamplerMixin:
"""Shared (epoch, position) persistence for the two resumable samplers."""
epoch: int
_position_in_epoch: int
def state_dict(self) -> dict[str, int]:
return {
"epoch": int(self.epoch),
"position_in_epoch": int(self._position_in_epoch),
}
def load_state_dict(self, state: dict[str, int]) -> None:
# Use set_epoch so any subclass-specific bookkeeping fires (e.g.
# DistributedSampler reseeds its shuffle generator off epoch).
self.set_epoch(int(state.get("epoch", 0)))
self._position_in_epoch = int(state.get("position_in_epoch", 0))
class ResumableRandomSampler(_ResumableSamplerMixin, Sampler[int]):
"""Single-rank shuffle sampler with deterministic per-epoch ordering.
Replaces the implicit RandomSampler inside DataLoader(shuffle=True) for
Phase 3 GRPO single-GPU training. Seeds the shuffle off (seed, epoch)
so the index sequence is reproducible across runs, and remembers how
many indices were yielded so resume picks up mid-epoch instead of
restarting from index 0. Drops the implicit non-determinism that
DataLoader(shuffle=True) ships with by default.
"""
def __init__(self, data_source, *, seed: int = 42, epoch: int = 0):
self.data_source = data_source
self.seed = int(seed)
self.epoch = int(epoch)
self._position_in_epoch = 0
def set_epoch(self, epoch: int) -> None:
self.epoch = int(epoch)
self._position_in_epoch = 0
def __iter__(self):
generator = torch.Generator()
generator.manual_seed(self.seed + self.epoch)
indices = torch.randperm(len(self.data_source), generator=generator).tolist()
skip = self._position_in_epoch
for offset, idx in enumerate(indices[skip:], start=skip):
# Update BEFORE yield: a Python generator suspended at `yield`
# never resumes if the consumer breaks. If we updated after the
# yield, an early break would leave the saved position pointing
# at the just-yielded index (so resume would re-yield it). The
# consumer is committed to consuming the value the moment
# __next__ returns, so charging position += 1 first matches
# "elements yielded" exactly.
self._position_in_epoch = offset + 1
yield idx
# Iter exhausted — caller should advance epoch + call set_epoch.
def __len__(self):
return len(self.data_source)
class ResumableDistributedSampler(_ResumableSamplerMixin, DistributedSampler):
"""DistributedSampler with mid-epoch resume support.
Index ordering matches the parent DistributedSampler exactly — the
Phase 3 contribution is only that we remember how many of those indices
have already been yielded and skip ahead on the next __iter__. Used by
Phase 3 GRPO under FSDP / DDP where the existing code already wires a
DistributedSampler via train_stage3_grpo.py.
"""
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self._position_in_epoch = 0
def set_epoch(self, epoch: int) -> None:
super().set_epoch(epoch)
self._position_in_epoch = 0
def __iter__(self):
# super().__iter__ yields a generator; materialise once so the
# skip-and-resume logic can index into the deterministic order.
base_indices = list(super().__iter__())
skip = self._position_in_epoch
for offset, idx in enumerate(base_indices[skip:], start=skip):
# See ResumableRandomSampler.__iter__: update position BEFORE
# yield to keep state consistent under early consumer-break.
self._position_in_epoch = offset + 1
yield idx