Skip to content

CLAUDE.md

Latest commit

 

History

History
457 lines (360 loc) · 16 KB

CLAUDE.md

File metadata and controls

457 lines (360 loc) · 16 KB
Raw

curator system

Overview

curator is a config-driven system for generating high-quality synthetic training datasets for post-training LLMs. The philosophy is to minimize code complexity through DataFrame-centric operations and configuration-driven workflows.

Core Philosophy

  • Config-Driven: All parameters come from YAML files, no hardcoding
  • DataFrame-Centric: Every operation uses inspectable DataFrames for debugging
  • Unified Loading: Single loader handles all formats (text, JSON, parquet)
  • Minimal Scripts: One unified processor replaces dozens of specialized scripts
  • Extraction-Based: Use LLMs to extract structured data from any source

Quick Start

# Run all rounds defined in config.yaml
python3 workflows/simple_unified_processor.py

# Run a specific round
python3 workflows/simple_unified_processor.py --round round_2 --config my_config.yaml

Core Architecture

Unified Processor (workflows/simple_unified_processor.py)

This is the primary workflow that supersedes all other processors. It handles any number of rounds based on configuration:

  • Automatically chains rounds (Round 2 uses Round 1 output, etc.)
  • All parameters from config.yaml (no hardcoding)
  • Supports validation, filtering, and structure preservation
  • Handles both text and parquet inputs seamlessly
  • Automatic directory organization: Creates config-based subdirectories (e.g., config_mcq.yaml → outputs/mcq/)

DataFrame-Centric Processing (processing/dataframe_core.py)

The entire system is built around DataFrame transformations, making every step inspectable with df.head():

Core DataFrame Pipeline Functions

  1. load_source_data_as_df(path, chunk_size=10000)

    • Input: Any file/directory (text, JSON, parquet, PDF, HTML)
    • Output: DataFrame with columns: [row_id, source_file, chunk_index, text, char_count]
    • For parquet: Preserves ALL existing columns from previous rounds

  2. apply_template_to_df(df, template_path)

    • Input: DataFrame with text/data columns
    • Output: Same DataFrame + prompt column
    • Uses .apply() to process each row through the template

  3. df_to_api_requests(df, model_name, max_tokens_per_part)

    • Input: DataFrame with prompt column
    • Output: JSONL files ready for batch API
    • Automatically splits into parts based on token limits
    • Only processes rows where is_valid != False

  4. results_to_df(jsonl_path)

    • Input: Extracted results from API (.extracted.jsonl)
    • Output: DataFrame with all extracted fields

  5. merge_results_with_source(source_df, results_df, has_branching=False)

    • Input: Original DataFrame + Results DataFrame
    • Output: Combined DataFrame preserving source text and metadata
    • has_branching=True: Use parent ID merge (for list extraction that creates new IDs)
    • has_branching=False: Use direct merge (for 1:1 transformations, default)
    • Handles compound row IDs (e.g., chunk_001_000)

Optional Enhancement Functions

  1. add_sampled_columns(df, variable_configs)

    • Adds random variables from YAML files (e.g., personas)
    • Example: df = add_sampled_columns(df, {'persona': 'variables/persona.yaml'})

  2. validate_required_columns(df, required_cols)

    • Adds/updates is_valid column based on null checks
    • Used before API submission to filter incomplete rows

  3. load_source_data_with_structure(root_dir)

    • Special loader that preserves subdirectory structure
    • Adds source_dir column and prefixes row_ids with directory name

Why DataFrames?

Every operation returns a DataFrame, enabling:

# Debug at any step
print(df.head())
print(df.columns)
print(df['prompt'].iloc[0])

# Quick filtering
valid_df = df[df['is_valid'] == True]

# Easy inspection
df.info()
df.describe()

This approach replaces hundreds of lines of custom processing logic with simple, composable DataFrame operations.

Configuration-Driven Pipeline

Example config.yaml

goal: Generate Q&A dataset from text
teacher_model: gemini-2.5-flash
api_type: vertex

input_data_metadata:
  path: data/             # Input directory or parquet file
  chunk_size: 100000      # Text chunking (ignored for parquet)
  tokens_per_part: 100000000  # Batch size for API

data_gen_procedures:
  round_1:
    prompt_name: multi_qa_extraction
    extraction_config: extraction_configs/multi_extraction.yaml
    has_branching: true  # List extraction creates new IDs
    
  round_2:
    # input_path: outputs/dataset/round1/  # Optional: load from split parquet
    prompt_name: solve_with_thinking
    extraction_config: extraction_configs/thinking_and_answer.yaml
    has_branching: false  # 1:1 transformation keeps same IDs
    validation:
      enabled: true
      strategy: exact_match
      compare_fields:
        predicted: model_answer
        ground_truth: answer
      filter: 'correct_only'

Parameter Hierarchy

Parameters cascade in this order:

  1. Round-specific config
  2. Top-level config (e.g., teacher_model)
  3. input_data_metadata defaults
  4. Hardcoded sensible defaults

Automatic Output Directory Organization

The system automatically creates config-based subdirectories to organize intermediate files and results:

# Config file determines output directory
config_mcq.yaml       →  outputs/mcq/
config_qa.yaml        →  outputs/qa/  
config_reasoning.yaml →  outputs/reasoning/

Benefits:

  • Clean separation: Each config gets its own directory
  • No conflicts: Multiple configs can run simultaneously without file collisions
  • Easy cleanup: Delete outputs/mcq/ to remove all files from that config
  • Organized results: Find all files related to a specific dataset in one place

Generated Files:

outputs/mcq/
├── round_1_requests.jsonl          # API batch requests
├── round_1_requests.metadata.jsonl # Request metadata
├── round_1_output.parquet          # Round results (auto-chained to Round 2)
├── round_2_requests.jsonl          # Next round requests
├── round_2_output.parquet          # Final results
└── source_with_structure.jsonl     # Directory structure info (if applicable)

Override behavior: You can still specify custom output_dir in any round config to use a different directory.

Key Features

Automatic Structure Preservation

When input has subdirectories, the system preserves structure through processing:

data/category_A/  →  outputs/dataset/category_A/
data/category_B/  →  outputs/dataset/category_B/

Unified Data Loading Philosophy

The key insight is that load_source_data_as_df() treats all input as text for LLM processing:

  • Text files: Chunks into specified size for context windows
  • JSON files: Loads as raw text - LLMs understand JSON structure naturally
  • Parquet files: Special case - preserves ALL columns from previous rounds, ignores chunk_size
  • PDFs/HTML: Extracts text content (requires optional libraries)

This unified approach means you don't need separate loaders or special handling - the LLM extracts what it needs based on your prompt template.

Validation System

Configure validation in any round:

validation:
  enabled: true
  strategy: exact_match  # or case_insensitive
  compare_fields:
    predicted: model_answer
    ground_truth: answer
  filter: all  # Keep all with is_valid column

Multi-Round Processing Concept

The system uses a pipeline approach where each round transforms data and passes it to the next:

  • Round 1: Extract structured data from raw text (e.g., Q&A pairs, MCQAs)
  • Round 2+: Apply transformations (e.g., solve questions, evaluate quality)
  • Final: Format for training (e.g., OpenAI chat format, DPO pairs)

Each round produces a parquet file that becomes input for the next round, maintaining full data lineage through compound row IDs (e.g., chunk_00000001_000_001).

Branching vs Non-Branching Rounds

Branching rounds (has_branching=True):

  • Create multiple outputs from single input (1:N transformation)
  • Example: Extract multiple Q&A pairs from one text chunk
  • IDs change: chunk_001 → chunk_001_000, chunk_001_001, chunk_001_002
  • Use parent ID merge when combining results

Non-branching rounds (has_branching=False, default):

  • Transform each input to single output (1:1 transformation)
  • Example: Solve an existing MCQA, evaluate quality
  • IDs stay same: chunk_001_000 → chunk_001_000
  • Use direct merge on row_id

Configure in merge function:

# Round 1: List extraction creates branching
merged_df = merge_results_with_source(source_df, results_df, has_branching=True)

# Round 2: Solving MCQAs, no new IDs created
merged_df = merge_results_with_source(source_df, results_df, has_branching=False)

Conceptual DataFrame Flow

# Round 1: Raw text → Questions
df = load_source_data_as_df('data/')          # Load raw text
df = apply_template_to_df(df, 'prompts/extract_qa.py')  # Add prompts
files = df_to_api_requests(df)                # Generate API requests
# ... API processing ...
results_df = results_to_df('results.jsonl')   # Load results
df = merge_results_with_source(df, results_df)  # Combine

# Round 2: Questions → Answers with thinking
df = load_source_data_as_df('round1.parquet') # Previous output
df = apply_template_to_df(df, 'prompts/solve.py')  # Solve prompts
# ... continue pipeline ...

Workflow

Complete Pipeline Example

# Complete Pipeline with Auto-Directory Detection
# All intermediate files automatically organized by config name

# Round 1: Extract MCQAs  
python3 workflows/simple_unified_processor.py --config config_mcq.yaml --round round_1
python3 io/upload_api_batch.py outputs/mcq/round_1_requests.jsonl --config config_mcq.yaml
python3 io/download_extract_results.py BATCH_JOB_NAME \
  --config config_mcq.yaml \
  --extraction-config extraction_configs/multi_extraction.yaml

# Post-process with automatic path detection
python3 workflows/post_process_with_structure.py \
  --config config_mcq.yaml \
  --extracted outputs/batch_results.extracted.jsonl \
  --source-parquet outputs/round_1_output.parquet \
  --has-branching true \
  --split-output

# Round 2: Solve MCQAs (auto-loads from config directory)
python3 workflows/simple_unified_processor.py --config config_mcq.yaml --round round_2  
python3 io/upload_api_batch.py outputs/mcq/round_2_requests.jsonl --config config_mcq.yaml
python3 io/download_extract_results.py BATCH_JOB_NAME \
  --config config_mcq.yaml \
  --extraction-config extraction_configs/thinking_and_answer.yaml

# Post-process with auto-detection and validation
python3 workflows/post_process_with_structure.py \
  --config config_mcq.yaml \
  --extracted outputs/batch_results.extracted.jsonl \
  --source-parquet outputs/round_2_output.parquet \
  --has-branching false \
  --round round_2 \
  --split-output

Filtering Data

When validation filters are specified in config.yaml (e.g., filter_1: model_answer == answer), you should apply these filters BEFORE running post_process_with_structure.py. The system will print a reminder about which filters to apply.

Example filters in config.yaml:

validation:
  enabled: true
  filter_1: model_answer == answer  # Keep only correct answers
  filter_2: thinking != null        # Remove rows with null thinking

When you see these filters, interpret them and apply using pandas:

import pandas as pd

# Load the extracted results
df = pd.read_json('outputs/batch_results.extracted.jsonl', lines=True)

# Apply filter_1: model_answer == answer
df = df[df['model_answer'] == df['answer']]

# Apply filter_2: thinking != null
df = df[df['thinking'].notna()]
df = df[df['thinking'] != '']  # Also remove empty strings

# Save filtered results
df.to_json('outputs/batch_results.extracted.filtered.jsonl', orient='records', lines=True)

Then use the filtered file in post_process_with_structure.py:

python3 workflows/post_process_with_structure.py \
  --extracted outputs/batch_results.extracted.filtered.jsonl \
  --source-parquet outputs/round_2_output.parquet \
  ...

Post-Processing with Structure (workflows/post_process_with_structure.py)

Handles post-API processing: merging results, applying validation, and preserving directory structure.

Key features:

  • Merges API results with source DataFrame
  • Applies validation filtering BEFORE splitting (if configured)
  • Preserves directory structure when splitting output
  • Overwrites source parquet for seamless round chaining

Usage:

# Basic usage
python3 workflows/post_process_with_structure.py \
  --extracted results.extracted.jsonl \
  --source-parquet outputs/round_N_output.parquet \
  --has-branching [true|false]

# With validation and structure
python3 workflows/post_process_with_structure.py \
  --extracted results.extracted.jsonl \
  --source-parquet outputs/round_2_output.parquet \
  --has-branching false \
  --round round_2 \              # Reads validation config
  --config config.yaml \          # From this file
  --split-output \                # Split by source_dir
  --output-dir outputs/filtered

Resuming from Split Parquet Files

The system automatically handles parquet files in subdirectories. To resume:

# config.yaml
round_2:
  input_path: outputs/dataset/round1/  # Loads all */processed.parquet files
  # ... rest of config

Then run normally:

python3 workflows/simple_unified_processor.py --round round_2

API Batch Processing

Upload & Download

# Upload batch
python3 io/upload_api_batch.py outputs/requests.jsonl

# Check status
python3 io/upload_api_batch.py --check-status BATCH_JOB_NAME

# Download and extract
python3 io/download_extract_results.py BATCH_JOB_NAME \
  --extraction-config extraction_configs/your_config.yaml

For Large Datasets (Multiple Parts)

# Process parts sequentially
python3 io/process_batch_parts_sequential.py \
  --part-dir outputs \
  --pattern "requests_part_*.jsonl"

Prompt Templates & Extraction

Template Philosophy

Prompts are Python files with f-string templates that reference DataFrame columns. This allows complex logic while keeping prompts readable and versionable. The extraction configs then parse the LLM's structured output.

Extraction Configs

Single Item Extraction

# extraction_configs/thinking_and_answer.yaml
patterns:
  thinking: '<THINKING>(.*?)</THINKING>'
  model_answer: '<ANSWER>\s*([A-D])\s*</ANSWER>'
defaults:
  thinking: null
  model_answer: null

List Extraction (Multiple Items)

# extraction_configs/multi_extraction.yaml
extraction_mode: list
item_pattern: '<ITEM>(.*?)</ITEM>'
patterns:
  question: '<QUESTION>(.*?)</QUESTION>'
  answer: '<ANSWER>(.*?)</ANSWER>'

Directory Structure

prompts/                  # Prompt templates (Python files)
extraction_configs/       # Regex patterns for result extraction (YAML)
format_templates/        # Output format definitions (Python)
cookbook/                # Example configurations
workflows/
  simple_unified_processor.py      # Main processor - generates API requests
  post_process_with_structure.py   # Post-processor - merges results, validates, splits
  simple_round1_processor.py       # Legacy - kept for reference
  simple_round2_processor.py       # Legacy - kept for reference
io/                      # API upload/download utilities
processing/
  dataframe_core.py      # Core DataFrame operations
  dataset_formatter.py   # Final formatting with templates

Important Notes

  • Always use simple_unified_processor.py - it supersedes all other processors
  • Validation defaults to enabled: false when not specified
  • The filter parameter controls whether invalid rows are removed immediately
  • For maximum flexibility, use filter: all and filter during final formatting
  • All JSON files are treated as plain text (the LLM extracts structured data)

Example Configurations

See cookbook/ directory for complete examples:

  • instruction_answer_generation_config.yaml: Two-round Q&A generation with evaluation
  • mcqa_generation_config.yaml: MCQA generation with thinking traces
  • config_unified_example.yaml: Comprehensive multi-round pipeline example