[tab] crowd-pilot/crowd-pilot/serialize_dataset_nemo_json.py (python)
#!/usr/bin/env python3
"""
CSV sessions -> JSONL (newline-delimited JSON) for NeMo 2.0 SFT function calling.

Each session is converted to a multi-turn conversation where:
- User role: terminal output (<stdout>...</stdout> content)
- Assistant role: terminal commands (bash commands)
"""

from __future__ import annotations

import argparse
import json
import os
from pathlib import Path
from typing import List, Tuple, cast
import random

import pandas as pd

from serialization_utils import (
    SerializeConfig,
    session_to_nemo_conversation,
    _discover_local_sessions,
)


def to_nemo_jsonl(cfg: SerializeConfig) -> None:
    """Convert CSV sessions to NeMo JSONL format."""
    os.makedirs(cfg.output_dir, exist_ok=True)

    required_cols = ["Sequence", "Time", "File", "RangeOffset", "RangeLength", "Text", "Language", "Type"]

    session_dataframes: List[Tuple[pd.DataFrame, str]] = []
    root = Path(cast(str, cfg.csv_root)).expanduser().resolve()
    csv_files = _discover_local_sessions(root)
    assert csv_files, f"No CSV files found under {root}"
    
    for csv_file in csv_files:
        df = pd.read_csv(csv_file)
        missing_local = [c for c in required_cols if c not in df.columns]
        assert not missing_local, f"Missing required CSV columns in {csv_file}: {missing_local}"
        session_dataframes.append((df, str(csv_file)))

    random.seed(42)
    random.shuffle(session_dataframes)
    
    total_sessions = len(session_dataframes)
    val_count = int(total_sessions * cfg.val_ratio)
    train_count = total_sessions - val_count

    train_conversations = []
    val_conversations = []
    
    session_turn_counts_all: List[int] = []
    session_turn_counts_kept: List[int] = []
    skipped_empty_conversations = 0
    docs_written = 0

    for i, (session_df, session_path) in enumerate(session_dataframes):
        if cfg.max_docs and docs_written >= cfg.max_docs:
            break
            
        session_df = pd.DataFrame(session_df.copy())
        
        # Convert directly to NeMo conversation format
        conversations = session_to_nemo_conversation(session_df)
        
        num_turns = len(conversations)
        session_turn_counts_all.append(num_turns)
        print(f"[debug] Session has {num_turns} conversation turns ({session_path})")

        # FIXME (f.srambical): skip sessions that are too short
        
        if not conversations:
            print(f"[warning] Session {session_path} produced no conversations")
            skipped_empty_conversations += 1
            continue
        
        record = {
            "mask": "User",
            "system": "",
            "conversations": conversations
        }
        
        if i < train_count:
            train_conversations.append(record)
        else:
            val_conversations.append(record)
        
        docs_written += 1

    train_path = Path(cfg.output_dir) / "train.jsonl"
    with open(train_path, 'w', encoding='utf-8') as f:
        for record in train_conversations:
            f.write(json.dumps(record, ensure_ascii=False) + '
')
    
    val_path = Path(cfg.output_dir) / "val.jsonl"
    with open(val_path, 'w', encoding='utf-8') as f:
        for record in val_conversations:
            f.write(json.dumps(record, ensure_ascii=False) + '
')

    def _print_turn_stats(label: str, values: List[int]) -> None:
        if not values:
            print(f"[debug] No {label} sessions for turn count stats.")
            return
        count = len(values)
        total = sum(values)
        avg = total / count
        min_len = min(values)
        max_len = max(values)
        print(
            f"[debug] {label.capitalize()} sessions turn stats: "
            f"count={count}, avg={avg:.1f}, min={min_len}, max={max_len}"
        )

    _print_turn_stats("all", session_turn_counts_all)
    _print_turn_stats("kept", session_turn_counts_kept)
    
    print(f"
[summary]")
    print(f"  Total sessions processed: {total_sessions}")
    print(f"  Sessions kept: {len(session_turn_counts_kept)}")
    print(f"  Skipped (too short): {len(session_turn_counts_all) - len(session_turn_counts_kept)}")
    print(f"  Skipped (empty conversations): {skipped_empty_conversations}")
    print(f"  Train conversations: {len(train_conversations)}")
    print(f"  Val conversations: {len(val_conversations)}")
    print(f"  Output: {cfg.output_dir}/{{train,val}}.jsonl")


def parse_args() -> SerializeConfig:
    p = argparse.ArgumentParser(
        description="Serialize CSV sessions to JSONL for NeMo 2.0 SFT function calling"
    )
    p.add_argument("--csv_root", type=str, required=True, 
                   help="Root directory containing per-session CSV files")
    p.add_argument("--output_dir", type=str, required=True, 
                   help="Output directory for JSONL files")
    p.add_argument("--min_session_chars", type=int, default=1024, 
                   help="Minimum characters to keep a session")
    p.add_argument("--max_docs", type=int, default=None, 
                   help="Stop after writing this many unique docs")
    p.add_argument("--val_ratio", type=float, default=0.10, 
                   help="Fraction of sessions to route to validation [0,1)")
    
    # These parameters are not used for JSONL output but kept for compatibility with SerializeConfig
    p.add_argument("--shard_size", type=int, default=20000, help="(unused for JSONL)")
    p.add_argument("--target_chars", type=int, default=8192, help="(unused for JSONL)")
    p.add_argument("--overlap_chars", type=int, default=128, help="(unused for JSONL)")
    p.add_argument("--arrayrecord_group_size", type=int, default=1, help="(unused for JSONL)")
    
    args = p.parse_args()
    return SerializeConfig(
        output_dir=args.output_dir,
        shard_size=args.shard_size,
        target_chars=args.target_chars,
        overlap_chars=args.overlap_chars,
        min_session_chars=args.min_session_chars,
        max_docs=args.max_docs,
        csv_root=(args.csv_root if args.csv_root else None),
        val_ratio=args.val_ratio,
        arrayrecord_group_size=args.arrayrecord_group_size,
    )


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


if __name__ == "__main__":
    main()