[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
import re
from pathlib import Path
from typing import List, Tuple, cast
import random

import pandas as pd

from serialization_utils import (
    SerializeConfig,
    session_to_bash_formatted_transcript,
    _discover_local_sessions,
)


def parse_transcript_to_conversation(transcript: str) -> List[dict]:
    """
    Parse a bash-formatted transcript into a list of conversation turns.
    
    Expected format alternates between:
    - bash code blocks (Assistant messages)
    - <stdout>...</stdout> blocks (User messages)
    
    Returns a list of {"from": "User"|"Assistant", "value": str} dictionaries.
    """
    conversations = []
    
    # Pattern to match bash code blocks: ```bash
...
```
    # Pattern to match stdout blocks: <stdout>
...
</stdout>
    
    # Split by alternating patterns
    # We'll process the transcript sequentially
    
    lines = transcript.split('
')
    i = 0
    
    while i < len(lines):
        line = lines[i]
        
        # Check for bash code block start
        if line.startswith('```bash'):
            # Extract bash command
            i += 1
            command_lines = []
            while i < len(lines) and not lines[i].startswith('```'):
                command_lines.append(lines[i])
                i += 1
            command = '
'.join(command_lines).strip()
            if command:
                conversations.append({
                    "from": "Assistant",
                    "value": command
                })
            i += 1  # Skip closing ```
            
        # Check for stdout block
        elif line.startswith('<stdout>'):
            # Extract stdout content
            i += 1
            output_lines = []
            while i < len(lines) and not lines[i].startswith('</stdout>'):
                output_lines.append(lines[i])
                i += 1
            output = '
'.join(output_lines).strip()
            if output:
                conversations.append({
                    "from": "User",
                    "value": f"<stdout>
{output}
</stdout>"
                })
            i += 1  # Skip closing </stdout>
            
        else:
            # Skip other content
            i += 1
    
    return conversations


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

    # Shuffle and split into train/val
    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 = []
    
    # Track statistics
    session_char_lengths_all: List[int] = []
    session_char_lengths_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())
        transcript = session_to_bash_formatted_transcript(session_df)
        
        session_chars = len(transcript.strip())
        session_char_lengths_all.append(session_chars)
        print(f"[debug] Session transcript length: {session_chars} chars ({session_path})")

        if session_chars < cfg.min_session_chars:
            print(f"Skipping session {session_path} because it's too short ({session_chars} chars)")
            continue
            
        session_char_lengths_kept.append(session_chars)
        
        # Parse transcript into conversations
        conversations = parse_transcript_to_conversation(transcript)
        
        if not conversations:
            print(f"[warning] Session {session_path} produced no conversations after parsing")
            skipped_empty_conversations += 1
            continue
        
        # Create NeMo format record
        record = {
            "mask": "User",
            "system": "",
            "conversations": conversations
        }
        
        # Route to train or val
        if i < train_count:
            train_conversations.append(record)
        else:
            val_conversations.append(record)
        
        docs_written += 1

    # Write train JSONL
    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) + '
')
    
    # Write val JSONL
    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) + '
')

    # Print statistics
    def _print_char_stats(label: str, values: List[int]) -> None:
        if not values:
            print(f"[debug] No {label} sessions for character length 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 char stats: "
            f"count={count}, avg={avg:.1f}, min={min_len}, max={max_len}"
        )

    _print_char_stats("all", session_char_lengths_all)
    _print_char_stats("kept", session_char_lengths_kept)
    
    print(f"
[summary]")
    print(f"  Total sessions processed: {total_sessions}")
    print(f"  Sessions kept: {len(session_char_lengths_kept)}")
    print(f"  Skipped (too short): {len(session_char_lengths_all) - len(session_char_lengths_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("--long_pause_threshold_ms", type=int, default=120000, 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,
        long_pause_threshold_ms=args.long_pause_threshold_ms,
        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()