import pandas as pd
import os

# Read the master CSV file
master_df = pd.read_csv('all_results.csv')

# Filter out rows where "success" is false
master_df = master_df[master_df['success']]

# Get unique model names
model_names = master_df['cfg_model_name'].unique()
model_names_short = [model_name.split('/')[-1] for model_name in model_names]
# keep unique model names
model_names = list(set(model_names_short))

# Create a directory to store the new CSV files
output_dir = 'model_csvs'
os.makedirs(output_dir, exist_ok=True)

# Specify the columns to keep
columns_to_keep = ['num_free_tokens', 'target_str', 'target_length', 'optimal_prompt', 'ratio', 'memorized']

# Iterate over each model name and create a new CSV file
for model_name in model_names:
    # Filter the DataFrame for the current model by checking if short name of config exactly matches
    # get the short name of the  master_df['cfg_model_name'] with lambda function
    model_df = master_df[master_df['cfg_model_name'].apply(lambda x: x.split('/')[-1] == model_name)]
    
    # Select the specified columns
    model_df = model_df[columns_to_keep]
    
    # Save the DataFrame to a new CSV file
    output_path = os.path.join(output_dir, f'{model_name}.csv')
    model_df.to_csv(output_path, index=False)

print("CSV files created successfully.")