| from __future__ import annotations |
|
|
| import hashlib |
| import math |
| import sys |
| from pathlib import Path |
| from typing import Any |
|
|
| import torch |
| from datasets import Dataset as HFDataset |
| from datasets import load_dataset |
| from torch.utils.data import BatchSampler, Dataset |
|
|
| SCRIPT_DIR = Path(__file__).resolve().parent |
| if str(SCRIPT_DIR) not in sys.path: |
| sys.path.insert(0, str(SCRIPT_DIR)) |
|
|
| from config import ensure_dir, fingerprint_payload |
|
|
|
|
| OBS_ROLE_NONE = 0 |
| OBS_ROLE_USER = 1 |
| OBS_ROLE_AGENT_FEEDBACK = 2 |
| SFT_CACHE_SCHEMA_VERSION = 2 |
|
|
| DURATION_BUCKET_ORDER = { |
| "short": 0, |
| "medium": 1, |
| "long": 2, |
| "hour_scale": 3, |
| } |
|
|
|
|
| def load_clean_split(dataset_cfg: dict[str, Any], split: str) -> HFDataset: |
| local_parquet_dir = dataset_cfg.get("cleaned_local_parquet_dir") |
| if local_parquet_dir and Path(local_parquet_dir).exists(): |
| parquet_path = Path(local_parquet_dir) / f"{split}.parquet" |
| if parquet_path.exists(): |
| return load_dataset("parquet", data_files={split: str(parquet_path)}, split=split) |
|
|
| cleaned_repo_id = dataset_cfg.get("cleaned_repo_id") |
| if not cleaned_repo_id: |
| raise ValueError("Dataset config must provide either cleaned_local_parquet_dir or cleaned_repo_id.") |
|
|
| source_split = dataset_cfg.get("source_split") |
| if source_split: |
| dataset = load_dataset(cleaned_repo_id, split=str(source_split)) |
| if split not in {"train", "validation"}: |
| return dataset |
|
|
| validation_ratio = float(dataset_cfg.get("validation_ratio", 0.02)) |
| split_seed = int(dataset_cfg.get("split_seed", 17)) |
| selected_indices = [ |
| index |
| for index, row in enumerate(dataset) |
| if assign_deterministic_split(str(row["id"]), validation_ratio, split_seed) == split |
| ] |
| return dataset.select(selected_indices) |
|
|
| return load_dataset(cleaned_repo_id, split=split) |
|
|
|
|
| def tokenize_text(tokenizer: Any, text: str | None, max_length: int) -> tuple[list[int], list[int]]: |
| if not text: |
| return [tokenizer.pad_token_id], [0] |
| encoded = tokenizer( |
| text, |
| add_special_tokens=True, |
| truncation=True, |
| max_length=max_length, |
| ) |
| return encoded["input_ids"], encoded["attention_mask"] |
|
|
|
|
| def tokenizer_supports_chat_template(tokenizer: Any) -> bool: |
| return bool(getattr(tokenizer, "chat_template", None)) |
|
|
|
|
| def should_use_base_chat_template(config: dict[str, Any], tokenizer: Any) -> bool: |
| dataset_cfg = config.get("dataset", {}) |
| inference_cfg = config.get("inference", {}) |
| if "use_base_chat_template" in inference_cfg: |
| return bool(inference_cfg["use_base_chat_template"]) and tokenizer_supports_chat_template(tokenizer) |
| return bool(dataset_cfg.get("use_base_chat_template", False)) and tokenizer_supports_chat_template(tokenizer) |
|
|
|
|
| def build_user_generation_observation(tokenizer: Any, text: str, use_base_chat_template: bool) -> str: |
| if use_base_chat_template: |
| return tokenizer.apply_chat_template( |
| [{"role": "user", "content": text}], |
| tokenize=False, |
| add_generation_prompt=True, |
| ) |
| return text |
|
|
|
|
| def build_assistant_feedback_observation(tokenizer: Any, text: str, use_base_chat_template: bool) -> str: |
| del tokenizer |
| if use_base_chat_template: |
| return f"<start_of_turn>model\n{text}<end_of_turn>\n" |
| return text |
|
|
|
|
| def build_assistant_decoder_target(tokenizer: Any, text: str, use_base_chat_template: bool) -> str: |
| if not use_base_chat_template: |
| return text |
|
|
| prompt_text = tokenizer.apply_chat_template( |
| [{"role": "user", "content": ""}], |
| tokenize=False, |
| add_generation_prompt=True, |
| ) |
| full_text = tokenizer.apply_chat_template( |
| [ |
| {"role": "user", "content": ""}, |
| {"role": "assistant", "content": text}, |
| ], |
| tokenize=False, |
| add_generation_prompt=False, |
| ) |
| if not full_text.startswith(prompt_text): |
| raise ValueError("Chat template full text did not start with the generation prompt prefix.") |
| return full_text[len(prompt_text) :] |
|
|
|
|
| def duration_bucket_rank(bucket: str | None) -> int: |
| return DURATION_BUCKET_ORDER.get(str(bucket), len(DURATION_BUCKET_ORDER)) |
|
|
|
|
| def assign_deterministic_split(row_id: str, validation_ratio: float, seed: int) -> str: |
| digest = hashlib.sha1(f"{seed}:{row_id}".encode("utf-8")).hexdigest() |
| score = int(digest[:8], 16) / 0xFFFFFFFF |
| return "validation" if score < validation_ratio else "train" |
|
|
|
|
| def stable_tick_index(time_seconds: float, tick_seconds: float) -> int: |
| ratio = max(0.0, float(time_seconds)) / tick_seconds |
| return int(math.ceil(ratio - 1e-9)) |
|
|
|
|
| def quantize_message_sequence( |
| messages: list[dict[str, Any]], |
| *, |
| tick_seconds: float, |
| feedback_delay_seconds: float, |
| ) -> list[dict[str, Any]]: |
| raw_events: list[dict[str, Any]] = [] |
| event_order = 0 |
| for message in messages: |
| speaker = str(message["speaker"]) |
| text = str(message["text"]) |
| time_seconds = float(message["t"]) |
|
|
| if speaker == "user": |
| raw_events.append( |
| { |
| "time_seconds": time_seconds, |
| "order": event_order, |
| "kind": "user", |
| "text": text, |
| } |
| ) |
| event_order += 1 |
| continue |
|
|
| raw_events.append( |
| { |
| "time_seconds": time_seconds, |
| "order": event_order, |
| "kind": "agent_speak", |
| "text": text, |
| } |
| ) |
| event_order += 1 |
| raw_events.append( |
| { |
| "time_seconds": time_seconds + feedback_delay_seconds, |
| "order": event_order, |
| "kind": "agent_feedback", |
| "text": text, |
| } |
| ) |
| event_order += 1 |
|
|
| raw_events.sort(key=lambda item: (float(item["time_seconds"]), int(item["order"]))) |
|
|
| quantized_events: list[dict[str, Any]] = [] |
| previous_tick = -1 |
| for event in raw_events: |
| desired_tick = stable_tick_index(float(event["time_seconds"]), tick_seconds) |
| assigned_tick = max(desired_tick, previous_tick + 1) |
| quantized_events.append( |
| { |
| "tick": assigned_tick, |
| "kind": event["kind"], |
| "text": event["text"], |
| } |
| ) |
| previous_tick = assigned_tick |
|
|
| return quantized_events |
|
|
|
|
| def resolve_bucket_horizon_ticks(duration_bucket: str, rollout_cfg: dict[str, Any]) -> int: |
| bucket_map = rollout_cfg.get("horizon_ticks_by_bucket", {}) |
| default_horizon = int(rollout_cfg.get("max_horizon_ticks", 36000)) |
| return int(bucket_map.get(duration_bucket, default_horizon)) |
|
|
|
|
| def build_fixed_tick_conversation( |
| *, |
| row: dict[str, Any], |
| tokenizer: Any, |
| config: dict[str, Any], |
| rollout_cfg: dict[str, Any], |
| max_observation_tokens: int, |
| max_decoder_tokens: int, |
| ) -> dict[str, Any] | None: |
| tick_seconds = float(rollout_cfg["tick_seconds"]) |
| chunk_ticks = int(rollout_cfg["chunk_ticks"]) |
| feedback_delay_seconds = float(rollout_cfg["post_speech_feedback_delay_seconds"]) |
| window_strategy = str(rollout_cfg.get("long_window_strategy", "tail")) |
| duration_bucket = str(row.get("meta", {}).get("duration_bucket", "short")) |
| use_base_chat_template = should_use_base_chat_template(config, tokenizer) |
|
|
| quantized_events = quantize_message_sequence( |
| messages=row["messages"], |
| tick_seconds=tick_seconds, |
| feedback_delay_seconds=feedback_delay_seconds, |
| ) |
| total_ticks = max((event["tick"] for event in quantized_events), default=0) + 1 |
|
|
| horizon_ticks = resolve_bucket_horizon_ticks(duration_bucket, rollout_cfg) |
| if bool(rollout_cfg.get("drop_overlong_examples", False)) and total_ticks > horizon_ticks: |
| return None |
|
|
| if total_ticks <= horizon_ticks: |
| effective_start_tick = 0 |
| effective_total_ticks = total_ticks |
| elif window_strategy == "tail": |
| effective_start_tick = total_ticks - horizon_ticks |
| effective_total_ticks = horizon_ticks |
| else: |
| effective_start_tick = 0 |
| effective_total_ticks = horizon_ticks |
|
|
| chunk_count = max(1, math.ceil(effective_total_ticks / chunk_ticks)) |
| chunk_lengths = [ |
| min(chunk_ticks, max(0, effective_total_ticks - chunk_index * chunk_ticks)) |
| for chunk_index in range(chunk_count) |
| ] |
| chunk_events: list[list[dict[str, Any]]] = [[] for _ in range(chunk_count)] |
|
|
| event_count = 0 |
| for event in quantized_events: |
| absolute_tick = int(event["tick"]) |
| if absolute_tick < effective_start_tick or absolute_tick >= effective_start_tick + effective_total_ticks: |
| continue |
|
|
| relative_tick = absolute_tick - effective_start_tick |
| chunk_index = relative_tick // chunk_ticks |
| offset = relative_tick % chunk_ticks |
|
|
| if event["kind"] == "user": |
| observation_role = OBS_ROLE_USER |
| observation_text = build_user_generation_observation( |
| tokenizer, |
| str(event["text"]), |
| use_base_chat_template, |
| ) |
| gate_target = 0 |
| decoder_text = None |
| elif event["kind"] == "agent_feedback": |
| observation_role = OBS_ROLE_AGENT_FEEDBACK |
| observation_text = build_assistant_feedback_observation( |
| tokenizer, |
| str(event["text"]), |
| use_base_chat_template, |
| ) |
| gate_target = 0 |
| decoder_text = None |
| else: |
| observation_role = OBS_ROLE_NONE |
| observation_text = None |
| gate_target = 1 |
| decoder_text = build_assistant_decoder_target( |
| tokenizer, |
| str(event["text"]), |
| use_base_chat_template, |
| ) |
|
|
| observation_input_ids, observation_attention_mask = tokenize_text( |
| tokenizer=tokenizer, |
| text=observation_text, |
| max_length=max_observation_tokens, |
| ) |
| decoder_labels, _ = tokenize_text( |
| tokenizer=tokenizer, |
| text=decoder_text, |
| max_length=max_decoder_tokens, |
| ) |
| if gate_target == 0: |
| decoder_labels = [] |
|
|
| chunk_events[chunk_index].append( |
| { |
| "offset": int(offset), |
| "absolute_tick": absolute_tick, |
| "observation_role": int(observation_role), |
| "observation_input_ids": observation_input_ids, |
| "observation_attention_mask": observation_attention_mask, |
| "gate_target": int(gate_target), |
| "decoder_labels": decoder_labels, |
| } |
| ) |
| event_count += 1 |
|
|
| return { |
| "row_id": str(row["id"]), |
| "scenario_id": row.get("meta", {}).get("scenario_id"), |
| "duration_bucket": duration_bucket, |
| "bucket_rank": duration_bucket_rank(duration_bucket), |
| "total_ticks": int(total_ticks), |
| "effective_start_tick": int(effective_start_tick), |
| "effective_total_ticks": int(effective_total_ticks), |
| "chunk_count": int(chunk_count), |
| "chunk_lengths": chunk_lengths, |
| "chunk_events": chunk_events, |
| "event_count": int(event_count), |
| } |
|
|
|
|
| def build_chunk_batch( |
| *, |
| conversations: list[dict[str, Any]], |
| chunk_index: int, |
| pad_token_id: int, |
| chunk_ticks: int, |
| tick_seconds: float, |
| ) -> dict[str, torch.Tensor] | None: |
| active_lengths: list[int] = [] |
| max_observation_tokens = 1 |
| max_decoder_tokens = 1 |
|
|
| for conversation in conversations: |
| if chunk_index >= int(conversation["chunk_count"]): |
| active_lengths.append(0) |
| continue |
| step_count = int(conversation["chunk_lengths"][chunk_index]) |
| active_lengths.append(step_count) |
| for event in conversation["chunk_events"][chunk_index]: |
| max_observation_tokens = max(max_observation_tokens, len(event["observation_input_ids"])) |
| max_decoder_tokens = max(max_decoder_tokens, len(event["decoder_labels"]) or 1) |
|
|
| max_steps = max(active_lengths, default=0) |
| if max_steps <= 0: |
| return None |
|
|
| batch_size = len(conversations) |
| observation_input_ids = torch.full( |
| (batch_size, max_steps, max_observation_tokens), |
| fill_value=pad_token_id, |
| dtype=torch.long, |
| ) |
| observation_attention_mask = torch.zeros((batch_size, max_steps, max_observation_tokens), dtype=torch.long) |
| decoder_labels = torch.full((batch_size, max_steps, max_decoder_tokens), fill_value=-100, dtype=torch.long) |
| tick_mask = torch.zeros((batch_size, max_steps), dtype=torch.bool) |
| gate_target = torch.zeros((batch_size, max_steps), dtype=torch.float32) |
| observation_role = torch.zeros((batch_size, max_steps), dtype=torch.long) |
| delta_seconds = torch.zeros((batch_size, max_steps), dtype=torch.float32) |
| elapsed_seconds = torch.zeros((batch_size, max_steps), dtype=torch.float32) |
|
|
| for batch_index, conversation in enumerate(conversations): |
| step_count = active_lengths[batch_index] |
| if step_count <= 0: |
| continue |
|
|
| tick_mask[batch_index, :step_count] = True |
| for local_tick in range(step_count): |
| global_tick = chunk_index * chunk_ticks + local_tick |
| delta_seconds[batch_index, local_tick] = 0.0 if global_tick == 0 else tick_seconds |
| elapsed_seconds[batch_index, local_tick] = global_tick * tick_seconds |
|
|
| for event in conversation["chunk_events"][chunk_index]: |
| offset = int(event["offset"]) |
| observation_role[batch_index, offset] = int(event["observation_role"]) |
| gate_target[batch_index, offset] = float(event["gate_target"]) |
|
|
| observation_ids = event["observation_input_ids"] |
| observation_mask = event["observation_attention_mask"] |
| observation_input_ids[batch_index, offset, : len(observation_ids)] = torch.tensor( |
| observation_ids, |
| dtype=torch.long, |
| ) |
| observation_attention_mask[batch_index, offset, : len(observation_mask)] = torch.tensor( |
| observation_mask, |
| dtype=torch.long, |
| ) |
|
|
| labels = event["decoder_labels"] |
| if labels: |
| decoder_labels[batch_index, offset, : len(labels)] = torch.tensor(labels, dtype=torch.long) |
|
|
| return { |
| "tick_mask": tick_mask, |
| "gate_target": gate_target, |
| "observation_role": observation_role, |
| "observation_input_ids": observation_input_ids, |
| "observation_attention_mask": observation_attention_mask, |
| "decoder_labels": decoder_labels, |
| "delta_seconds": delta_seconds, |
| "elapsed_seconds": elapsed_seconds, |
| } |
|
|
|
|
| class ThoughtLoopConversationDataset(Dataset): |
| def __init__(self, config: dict[str, Any], tokenizer: Any, split: str) -> None: |
| self.config = config |
| self.tokenizer = tokenizer |
| self.split = split |
| self.examples = self._load_or_build() |
|
|
| def _cache_path(self) -> Path: |
| dataset_cfg = self.config["dataset"] |
| model_cfg = self.config["model"] |
| rollout_cfg = self.config["rollout"] |
| cache_cfg = self.config["cache"] |
| cache_root = ensure_dir(cache_cfg["preprocessed_root"]) |
| payload = { |
| "cache_schema_version": SFT_CACHE_SCHEMA_VERSION, |
| "split": self.split, |
| "dataset": dataset_cfg, |
| "model_name": model_cfg["base_model_name"], |
| "rollout": rollout_cfg, |
| "max_observation_tokens": model_cfg["max_observation_tokens"], |
| "max_decoder_tokens": model_cfg["max_decoder_tokens"], |
| "tokenizer_vocab_size": getattr(self.tokenizer, "vocab_size", None), |
| } |
| return cache_root / f"{self.split}_{fingerprint_payload(payload)}.pt" |
|
|
| def _load_or_build(self) -> list[dict[str, Any]]: |
| cache_path = self._cache_path() |
| if cache_path.exists(): |
| return torch.load(cache_path, map_location="cpu") |
|
|
| dataset = load_clean_split(self.config["dataset"], self.split) |
| rollout_cfg = self.config["rollout"] |
| model_cfg = self.config["model"] |
|
|
| examples: list[dict[str, Any]] = [] |
| for row in dataset: |
| example = build_fixed_tick_conversation( |
| row=row, |
| tokenizer=self.tokenizer, |
| config=self.config, |
| rollout_cfg=rollout_cfg, |
| max_observation_tokens=int(model_cfg["max_observation_tokens"]), |
| max_decoder_tokens=int(model_cfg["max_decoder_tokens"]), |
| ) |
| if example is not None: |
| examples.append(example) |
|
|
| if bool(rollout_cfg.get("sort_by_duration_bucket", True)): |
| examples.sort(key=lambda item: (int(item["bucket_rank"]), int(item["total_ticks"]), str(item["row_id"]))) |
|
|
| torch.save(examples, cache_path) |
| return examples |
|
|
| def __len__(self) -> int: |
| return len(self.examples) |
|
|
| def __getitem__(self, index: int) -> dict[str, Any]: |
| return self.examples[index] |
|
|
|
|
| class DurationBucketBatchSampler(BatchSampler): |
| def __init__( |
| self, |
| dataset: ThoughtLoopConversationDataset, |
| batch_size: int, |
| ) -> None: |
| self.dataset = dataset |
| self.batch_size = batch_size |
| self.batches: list[list[int]] = [] |
|
|
| current_bucket: str | None = None |
| current_batch: list[int] = [] |
| for index, example in enumerate(self.dataset.examples): |
| bucket = str(example["duration_bucket"]) |
| if current_batch and (bucket != current_bucket or len(current_batch) >= self.batch_size): |
| self.batches.append(current_batch) |
| current_batch = [] |
| current_bucket = bucket |
| current_batch.append(index) |
| if len(current_batch) >= self.batch_size: |
| self.batches.append(current_batch) |
| current_batch = [] |
|
|
| if current_batch: |
| self.batches.append(current_batch) |
|
|
| def __iter__(self) -> Any: |
| return iter(self.batches) |
|
|
| def __len__(self) -> int: |
| return len(self.batches) |
|
|
|
|
| def identity_collate(batch: list[dict[str, Any]]) -> list[dict[str, Any]]: |
| return batch |
|
|
|
|
| def estimate_total_chunk_microsteps( |
| *, |
| dataset: ThoughtLoopConversationDataset, |
| batch_sampler: DurationBucketBatchSampler, |
| ) -> int: |
| total_microsteps = 0 |
| for batch_indices in batch_sampler.batches: |
| total_microsteps += max(int(dataset.examples[index]["chunk_count"]) for index in batch_indices) |
| return total_microsteps |
|
|
