[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 hashlib
import os
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable, List, Optional, Set, Tuple, cast

import pandas as pd
from datasets import Dataset, DatasetDict, Features, Value, load_dataset, concatenate_datasets



@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]


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 _session_to_documents(df: pd.DataFrame, target_chars: int, overlap_chars: int, min_doc_chars: int) -> List[str]:
    """
    Convert a single session (CSV) into one or more training documents.
    Strategy: gather latest content per file, then assemble a session doc and chunk.
    """

    # Latest content per file by Sequence (assumed monotonic per CSV); groupby File and take last row
    df_sorted = df.sort_values(["File", "Sequence"])  # ensure order within file
    latest_per_file = df_sorted.groupby("File", as_index=False).tail(1)

    # Build a session-level document with multiple files
    parts: List[str] = []
    latest_map: dict[str, str] = {}
    for _, row in latest_per_file.iterrows():
        file_path = str(row.get("File", ""))
        lang_val = row.get("Language", "")
        lang = lang_val if isinstance(lang_val, str) else ""
        text = _clean_text(str(row.get("Text", "")))
        if not text:
            continue
        # Drop extremely small snippets to avoid noise
        if len(text) < 16:
            continue
        parts.append(_fenced_block(file_path, lang, text))
        latest_map[file_path] = text

    # Append selection events explicitly so the LM can learn cursor/mouse behavior.
    # We attempt to extract the selected substring from the latest snapshot when possible; otherwise we log metadata only.
    selection_rows = df.sort_values(["Sequence"]).query("Type == 'selection_command'")
    selection_parts: List[str] = []
    for _, srow in selection_rows.iterrows():
        file_path = str(srow.get("File", ""))
        lang_val = srow.get("Language", "")
        lang = lang_val if isinstance(lang_val, str) else ""
        ro = srow.get("RangeOffset", 0)
        rl = srow.get("RangeLength", 0)
        offset = int(ro) if isinstance(ro, (int, float, str)) and str(ro).isdigit() else 0
        length = int(rl) if isinstance(rl, (int, float, str)) and str(rl).isdigit() else 0
        header = f"// selection: file={file_path} offset={offset} length={length}"
        snippet = ""
        base_text = latest_map.get(file_path)
        if base_text and length > 0 and 0 <= offset < len(base_text):
            end = min(offset + length, len(base_text))
            snippet = base_text[offset:end]
        if snippet:
            selection_parts.append(header + "
" + _fenced_block(file_path, lang, snippet))
        else:
            selection_parts.append(header)

    if not parts:
        return []

    session_doc = "

".join(parts).strip()
    if selection_parts:
        session_doc = session_doc + "

// selections
" + "

".join(selection_parts)
    if len(session_doc) < min_doc_chars:
        # Small sessions can still be useful; return single doc if not empty
        return [session_doc] if session_doc else []

    # Chunk by characters with overlap
    docs: List[str] = []
    start = 0
    n = len(session_doc)
    while start < n:
        end = min(start + target_chars, n)
        chunk = session_doc[start:end].strip()
        if len(chunk) >= min_doc_chars:
            docs.append(chunk)
        if end == n:
            break
        start = max(end - overlap_chars, start + 1)
    return docs


def _hash(text: str) -> str:
    return hashlib.sha1(text.encode("utf-8")).hexdigest()


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", "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]

    seen: Set[str] = set()
    rows_out = 0
    shard_idx = 0
    buffer: List[str] = []

    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())
        breakpoint()
        docs = _session_to_documents(
            session_df,
            target_chars=cfg.target_chars,
            overlap_chars=cfg.overlap_chars,
            min_doc_chars=cfg.min_doc_chars,
        )
        for d in docs:
            h = _hash(d)
            if h in seen:
                continue
            seen.add(h)
            buffer.append(d)
            if cfg.max_docs and len(seen) >= cfg.max_docs:
                break
            if len(buffer) >= cfg.shard_size:
                flush_buffer(shard_idx)
                shard_idx += 1
        if cfg.max_docs and len(seen) >= 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")
    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,
    )


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


if __name__ == "__main__":
    main()