reformat_csv.py

#!/usr/bin/env python3
"""
Reformat a CSV file from the multi-LLM-per-row format (like sap_facility_custom_routing_merged.csv)
to the one-LLM-per-row format (like router_arena_data/test2.csv).

Source format:
  - Columns: prompt_id, prompt_id, prompt, {model}/response, {model}/score, ...
  - One row per prompt with multiple model response/score columns

Target format:
  - Columns: id, original_id, question, ground_truth, completion, input_tokens, output_tokens, cost, performance, task, llm
  - One row per model response, with original_id linking rows from the same prompt
"""

import argparse
import csv
import re
import sys
from pathlib import Path

csv.field_size_limit(sys.maxsize)

def extract_model_columns(headers: list[str]) -> list[tuple[str, str, str]]:
    """
    Extract model names and their response/score column indices.
    
    Returns list of tuples: (model_name, response_col, score_col)
    """
    models = []
    response_pattern = re.compile(r'^(.+)/response$')
    
    for header in headers:
        match = response_pattern.match(header)
        if match:
            model_name = match.group(1)
            score_col = f"{model_name}/score"
            if score_col in headers:
                models.append((model_name, header, score_col))
    
    return models


def reformat_csv(input_file: str, output_file: str) -> None:
    """
    Reformat CSV from multi-LLM-per-row to one-LLM-per-row format.
    """
    input_path = Path(input_file)
    output_path = Path(output_file)
    
    if not input_path.exists():
        print(f"Error: Input file '{input_file}' not found.")
        sys.exit(1)
    
    with open(input_path, 'r', newline='', encoding='utf-8') as infile:
        reader = csv.DictReader(infile)
        headers = reader.fieldnames
        
        if not headers:
            print("Error: Could not read headers from input file.")
            sys.exit(1)
        
        # Find the prompt column (could be 'prompt' or 'question')
        prompt_col = None
        for col in ['prompt', 'question']:
            if col in headers:
                prompt_col = col
                break
        
        if not prompt_col:
            print("Error: Could not find 'prompt' or 'question' column in input file.")
            sys.exit(1)
        
        # Extract model columns
        models = extract_model_columns(headers)
        
        if not models:
            print("Error: Could not find any model response/score column pairs.")
            print(f"Headers found: {headers}")
            sys.exit(1)
        
        print(f"Found {len(models)} models:")
        for model_name, _, _ in models:
            print(f"  - {model_name}")
        
        # Output columns matching target format
        output_headers = [
            'id', 'original_id', 'question', 'ground_truth', 'completion',
            'input_tokens', 'output_tokens', 'cost', 'performance', 'task', 'llm'
        ]
        
        with open(output_path, 'w', newline='', encoding='utf-8') as outfile:
            writer = csv.DictWriter(outfile, fieldnames=output_headers)
            writer.writeheader()
            
            row_id = 0
            original_id = 0
            
            for row in reader:
                prompt = row.get(prompt_col, '')
                
                # Create one output row per model
                for model_name, response_col, score_col in models:
                    response = row.get(response_col, '')
                    score = row.get(score_col, '')
                    
                    output_row = {
                        'id': row_id,
                        'original_id': original_id,
                        'question': prompt,
                        'ground_truth': 'none',
                        'completion': response,
                        'input_tokens': 'none',
                        'output_tokens': 'none',
                        'cost': 'none',
                        'performance': score,
                        'task': 'none',
                        'llm': model_name
                    }
                    writer.writerow(output_row)
                    row_id += 1
                
                original_id += 1
        
        print(f"\nReformatted {original_id} prompts into {row_id} rows.")
        print(f"Output written to: {output_path}")


def main():
    parser = argparse.ArgumentParser(
        description='Reformat CSV from multi-LLM-per-row to one-LLM-per-row format.',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog='''
Examples:
  python reformat_csv.py input.csv output.csv
  python reformat_csv.py sap_facility_data/sap_facility_custom_routing_merged.csv reformatted_output.csv
        '''
    )
    parser.add_argument('--input-file', help='Input CSV file (multi-LLM-per-row format)')
    parser.add_argument('--output-file', help='Output CSV file (one-LLM-per-row format)')
    
    args = parser.parse_args()
    
    reformat_csv(args.input_file, args.output_file)


if __name__ == '__main__':
    main()