[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


# Heuristic sets. Adjust as needed for your capture.
CODE_LANGS = {
    "python", "py", "typescript", "ts", "javascript", "js", "tsx", "jsx",
    "go", "rust", "rs", "java", "kotlin", "kt", "scala", "c", "h",
    "cpp", "cc", "hpp", "cxx", "cs", "swift", "ruby", "rb", "php",
    "sh", "bash", "zsh", "r", "jl", "sql", "cuda", "cu", "metal",
    "yaml", "yml", "toml", "ini", "json", "xml", "html", "css", "scss",
    "md", "markdown", "tex",
}

# Rows with Type that are clearly not training-signal; keep "content" and "tab" but drop purely telemetry.
SKIP_TYPES = {"selection_command"}


@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 _is_codey_language(lang: Optional[str]) -> bool:
    if not lang:
        return False
    lang = str(lang).lower()
    return lang in CODE_LANGS


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()
    # Prefer the declared language if code-like; otherwise leave empty fence for tokenizer neutrality.
    fence_lang = lang if _is_codey_language(lang) else ""
    return f"// file: {path}
```{fence_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] = []
    for _, row in latest_per_file.iterrows():
        file_path = str(row.get("File", ""))
        lang_val = row.get("Language", "")
        # Robust NA handling without triggering type checkers
        try:
            is_na = pd.isna(lang_val)
        except Exception:
            is_na = False
        lang = "" if is_na else str(lang_val)
        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))

    if not parts:
        return []

    session_doc = "

".join(parts).strip()
    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:
        # Treat hf_path as a dataset repository id (recommended). For raw CSV files, use --data_files.
        loaded = load_dataset(hf_path, split=split)
    elif data_files:
        loaded = load_dataset("csv", data_files=data_files, split=split)
    else:
        raise ValueError("Provide either --hf_path or --data_files")

    if isinstance(loaded, Dataset):
        return loaded
    if isinstance(loaded, DatasetDict):
        if split in loaded:
            return cast(Dataset, loaded[split])
        # Fallback: take first
        first_key = next(iter(loaded.keys()))
        return cast(Dataset, loaded[first_key])
    # IterableDataset variants not expected here; coerce via to_iterable_dataset if needed.
    raise TypeError("Expected a Dataset or DatasetDict from load_dataset")


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]
    if missing:
        raise ValueError(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.
    session_key = "SessionId" if "SessionId" in ds.column_names else None

    # Convert to pandas in partitions to keep memory bounded.
    # If the dataset is large, you may prefer streaming + manual grouping. Here we iterate in HF shards.
    shards: List[Dataset] = []
    if isinstance(ds, Dataset):
        shards = [ds]
    elif isinstance(ds, DatasetDict):
        # defensive: some loaders produce DatasetDict
        shards = [cast(Dataset, ds[k]) for k in ds.keys()]
    else:
        shards = [ds]  # type: ignore[list-item]

    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:
        # to_dataframe reads whole shard in-memory. Consider .iter(batch_size=...) for very large datasets.
        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 = cast(pd.DataFrame, session_df_obj)
            session_df = session_df[session_df["Language"].map(_is_codey_language)]
            if session_df.empty:
                continue
            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()