[content] 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
from tqdm import tqdm


# Event types to skip entirely (keep empty to preserve all; we model selections explicitly)
SKIP_TYPES: set[str] = set()


@dataclass
class SerializeConfig:
    hf_path: Optional[str]
    split: str
    data_files: Optional[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.
    """
    # Keep only meaningful rows
    keep = (~df["Type"].isin(list(SKIP_TYPES))) & df["Text"].notna() & (df["Text"].str.len() > 0)
    df = df.loc[keep]

    # 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 ""
        try:
            offset = int(srow.get("RangeOffset", 0))
        except Exception:
            offset = 0
        try:
            length = int(srow.get("RangeLength", 0))
        except Exception:
            length = 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: Optional[str], split: str, data_files: Optional[str]) -> Dataset:
    if hf_path:
        loaded = load_dataset(hf_path, split=split)
    elif data_files:
        # NOTE (f.srambical): this is solely for local testing purposes
        loaded = load_dataset("csv", data_files=data_files, split=split)
    else:
        raise ValueError("Provide either --hf_path or --data_files")

    assert isinstance(loaded, DatasetDict), "Expected a DatasetDict from load_dataset"
    assert split in loaded, f"Split {split} not found in dataset"
    return cast(Dataset, loaded[split])


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

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

    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}"

    # We will group by session if present; otherwise treat whole split as one meta-session with per-file grouping.
    # FIXME (f.srambical): this is not the case when loading a local CSV. Check whether this is the case when loading from huggingface
    session_key = "SessionId" if "SessionId" in ds.column_names else None

    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())

        if session_key and session_key in pdf.columns:
            grouped = pdf.groupby(session_key)
            groups_iter = grouped
        else:
            # Single pseudo-session to drive grouping by file
            pdf = pdf.copy()
            pdf["__PseudoSession__"] = 0
            groups_iter = pdf.groupby("__PseudoSession__")

        for _, session_df_obj in tqdm(groups_iter, desc="sessions"):
            # Within a session, filter languages early
            session_df = pd.DataFrame(session_df_obj)
            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, default=None, help="HF repo id for the dataset (uses csv builder)")
    p.add_argument("--data_files", type=str, default=None, help="Local CSV glob or mapping understood by datasets.load_dataset")
    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,
        data_files=args.data_files,
        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()