[tab] crowd-pilot/serialize_dataset.py (python)
#!/usr/bin/env python3
"""
CSV sessions -> Parquet shards (single 'text' column) for MaxText Grain pretraining.

Design:
- Load a Hugging Face dataset pointing to CSV files (local or hub). We do not stream to allow grouping.
- Each CSV row describes an IDE event with code snapshot in the 'Text' field and metadata columns.
- For pretraining, we extract language-modelable text and produce a corpus of documents under a size cap.

Transform steps:
1) Column normalization and filtering:
   - Keep only rows where 'Language' looks like code or doctext (e.g., python, typescript, markdown, plaintext).
   - Optionally skip rows with Type that are pure logs/telemetry.
2) Group rows by 'File' within a session and concatenate contiguous edits into file-level snapshots.
   - These represent evolving code; we prefer final or merged snapshots to reduce near-duplicates.
3) Build session-level documents by concatenating a small preamble with file path and fenced code blocks, in source order.
   - We include lightweight separators between files to help the tokenizer learn boundaries without heavy prompt tokens.
4) Chunking:
   - Token-agnostic length-aware chunking (by characters) to target a safe upper bound that fits `max_target_length`.
   - Defaults: target_chars ~= 8192, overlap ~= 128 chars to keep continuity; adjust by CLI.
   - MaxText Grain will still tokenize and pack, but keeping documents reasonable improves throughput.
5) Deduplication:
   - Hash-based document dedup across all sessions/files.
   - Optional line-based minhash could be added later; a simple SHA1 is sufficient initially.
6) Write Parquet shards with a single column 'text'.

Usage:
  python -m crowd_pilot.serialize_dataset \
    --hf_path your-org/your-dataset  \
    --split train \
    --output_dir ./out_parquet \
    --shard_size 20000 \
    --target_chars 8192 \
    --overlap_chars 128

You can also use --data_files to point to explicit CSV files.

Downstream MaxText config (examples):
  dataset_type: grain
  grain_file_type: parquet
  grain_train_files: '/abs/path/out_parquet/*.parquet'
  train_data_columns: ['text']
  tokenize_train_data: True

Notes:
- We keep only 'text' column per Grain pipeline expectations for pretraining.
- No tokenization here; MaxText handles tokenization/packing.
"""

from __future__ import annotations

import argparse
import os
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable, List, Optional, Set, Tuple, cast, Dict

import pandas as pd  # type: ignore
from datasets import Dataset, DatasetDict, Features, Value, load_dataset, concatenate_datasets  # type: ignore



@dataclass
class SerializeConfig:
    hf_path: str
    split: str
    output_dir: str
    shard_size: int
    target_chars: int
    overlap_chars: int
    min_doc_chars: int
    max_docs: Optional[int]
    keyframe_every_events: int
    long_pause_threshold_ms: int


def _clean_text(text: str) -> str:
    # Normalize line endings and strip trailing spaces; preserve tabs/newlines.
    return text.replace("
", "
").replace("
", "
").rstrip()


def _fenced_block(path: str, language: Optional[str], content: str) -> str:
    lang = (language or "").lower()
    return f"// file: {path}
```{lang}
{content}
```
"


def _apply_change(content: str, offset: int, length: int, new_text: str) -> str:
    # Mirrors crowd_code_player.replay_file.apply_change
    base = str(content)
    text = str(new_text) if pd.notna(new_text) else ""
    text = text.replace("\
", "
").replace("\
", "
")
    if offset > len(base):
        base = base + (" " * (offset - len(base)))
    return base[:offset] + text + base[offset + length:]


def _session_to_transcript(
    df: pd.DataFrame,
    keyframe_every_events: int,
    long_pause_threshold_ms: int,
) -> str:
    # Prefer Time if present, else Sequence
    breakpoint()

    file_states: Dict[str, str] = {}
    terminal_state: str = ""
    active_file: Optional[str] = None
    per_file_event_counts: Dict[str, int] = {}
    last_time_ms: Optional[int] = None

    parts: List[str] = []

    def emit_keyframe_for(file_path: str, language: Optional[str]) -> None:
        if file_path == "TERMINAL":
            if not terminal_state:
                return
            parts.append("// observation: terminal")
            parts.append(_fenced_block("TERMINAL", None, _clean_text(terminal_state)))
        else:
            # FIXME (f.srambical): this yields the entire file content, but we likely only want to emit the changes, or a local viewport
            content = file_states.get(file_path, "")
            if not content:
                return
            lang = (language or "") if isinstance(language, str) else ""
            parts.append(f"// observation: file={file_path}")
            parts.append(_fenced_block(file_path, lang, _clean_text(content)))

    for i in range(len(df)):
        row = df.iloc[i]
        file_path: str = row["File"]
        language: Optional[str] = row["Language"]
        event_time: int = row["Time"]

        # Long pause detection
        if last_time_ms is not None:
            delta = event_time - last_time_ms
            if delta > long_pause_threshold_ms:
                # TODO (f.srambical): think about whether we want to emit this as an observation or not
                parts.append(f"<obs long_pause ms=\"{delta}\" />")
                if active_file is not None:
                    emit_keyframe_for(active_file, language)
        last_time_ms = event_time

        # Focus / switch
        if active_file != file_path:
            active_file = file_path
            if active_file == "TERMINAL":
                parts.append("<act focus target=\"terminal\" />")
            else:
                parts.append(f"<act focus file=\"{active_file}\" />")
            emit_keyframe_for(active_file, language)

        # Terminal input/output events
        if file_path == "TERMINAL":
            raw_text = row.get("Text", "")
            term_fragment = str(raw_text) if pd.notna(raw_text) else ""
            term_fragment = term_fragment.replace("\
", "
").replace("\
", "
")
            parts.append("<act terminal_input />")
            prev_len = len(terminal_state)
            # Append plus newline as in the replayer
            terminal_state = terminal_state + term_fragment + "
"
            appended = terminal_state[prev_len:]
            if appended.strip():
                parts.append("// observation: terminal_delta")
                parts.append(_fenced_block("TERMINAL", None, _clean_text(appended)))
            continue

        # File edit events
        ro = row.get("RangeOffset", 0)
        rl = row.get("RangeLength", 0)
        new_text = row.get("Text", "")
        offset = _safe_int(ro, 0)
        length = _safe_int(rl, 0)
        new_text_str = str(new_text) if pd.notna(new_text) else ""

        operation = "noop"
        if length == 0 and new_text_str:
            operation = "insert"
        elif length > 0 and not new_text_str:
            operation = "delete"
        elif length > 0 and new_text_str:
            operation = "replace"

        parts.append(f"<act {operation} file=\"{file_path}\" offset=\"{offset}\" len=\"{length}\" />")

        before = file_states.get(file_path, "")
        after = _apply_change(before, offset, length, new_text)
        file_states[file_path] = after
        per_file_event_counts[file_path] = per_file_event_counts.get(file_path, 0) + 1

        if per_file_event_counts[file_path] % max(1, keyframe_every_events) == 0:
            emit_keyframe_for(file_path, language)

        # Selection/move events (optional tokenization)
        evt_type = str(row.get("Type", ""))
        if evt_type == "selection_command":
            parts.append(f"<act select file=\"{file_path}\" offset=\"{offset}\" len=\"{length}\" />")

    # Final keyframe for active file to end the transcript
    if active_file is not None:
        emit_keyframe_for(active_file, None)

    return "
".join(parts).strip()


 


def load_hf_csv(hf_path: str, split: str) -> Dataset:
    loaded = load_dataset(hf_path, split=split)

    assert isinstance(loaded, Dataset), "Expected a Dataset from load_dataset"
    return loaded


def to_parquet(
    cfg: SerializeConfig,
) -> None:
    os.makedirs(cfg.output_dir, exist_ok=True)

    ds = load_hf_csv(cfg.hf_path, cfg.split)

    required_cols = ["Sequence", "Time", "File", "RangeOffset", "RangeLength", "Text", "Language", "Type"]
    missing = [c for c in required_cols if c not in ds.column_names]
    assert not missing, f"Missing required CSV columns: {missing}"

    shards: List[Dataset] = [ds]

    rows_out = 0
    shard_idx = 0
    buffer: List[str] = []
    docs_written = 0

    def flush_buffer(idx: int):
        nonlocal buffer, rows_out
        if not buffer:
            return
        df_out = pd.DataFrame({"text": buffer})
        out_path = Path(cfg.output_dir) / f"shard_{idx:05d}.parquet"
        table = df_out
        table.to_parquet(out_path, index=False)
        rows_out += len(df_out)
        buffer = []

    for shard in shards:
        pdf = cast(pd.DataFrame, shard.to_pandas())
        session_df = pd.DataFrame(pdf.copy())
        transcript = _session_to_transcript(
            session_df,
            keyframe_every_events=cfg.keyframe_every_events,
            long_pause_threshold_ms=cfg.long_pause_threshold_ms,
        )
        buffer.append(transcript)
        docs_written += 1
        if len(buffer) >= cfg.shard_size:
            flush_buffer(shard_idx)
            shard_idx += 1
        if cfg.max_docs and docs_written >= cfg.max_docs:
            break

    if buffer:
        flush_buffer(shard_idx)

    print(f"Wrote {rows_out} documents to {cfg.output_dir}")


def parse_args() -> SerializeConfig:
    p = argparse.ArgumentParser(description="Serialize HF CSV sessions to Parquet for MaxText Grain")
    p.add_argument("--hf_path", type=str, required=True, help="HF repo id for the dataset (uses csv builder)")
    p.add_argument("--split", type=str, default="train", help="Split name to load")
    p.add_argument("--output_dir", type=str, required=True, help="Output directory for Parquet shards")
    p.add_argument("--shard_size", type=int, default=20000, help="Rows per Parquet shard")
    p.add_argument("--target_chars", type=int, default=8192, help="Target characters per document chunk")
    p.add_argument("--overlap_chars", type=int, default=128, help="Character overlap between chunks")
    p.add_argument("--min_doc_chars", type=int, default=256, help="Minimum characters to keep a chunk")
    p.add_argument("--max_docs", type=int, default=None, help="Stop after writing this many unique docs")
    p.add_argument("--keyframe_every_events", type=int, default=40, help="Emit an observation snapshot every N edit events per file")
    p.add_argument("--long_pause_threshold_ms", type=int, default=120000, help="Threshold (ms) to annotate long pauses and emit a keyframe")
    args = p.parse_args()
    return SerializeConfig(
        hf_path=args.hf_path,
        split=args.split,
        output_dir=args.output_dir,
        shard_size=args.shard_size,
        target_chars=args.target_chars,
        overlap_chars=args.overlap_chars,
        min_doc_chars=args.min_doc_chars,
        max_docs=args.max_docs,
        keyframe_every_events=args.keyframe_every_events,
        long_pause_threshold_ms=args.long_pause_threshold_ms,
    )


def main() -> None:
    cfg = parse_args()
    to_parquet(cfg)


if __name__ == "__main__":
    main()