import pandas as pd

def filter_dataset(dataset_path, model_name='llama-v3.3-70b-instruct'):
    """
    Filter dataset to only include statements with a label
    """
    dataset = pd.read_csv(dataset_path)
    dataset_path = dataset_path.split('.')[0]
    model_dataset = dataset[dataset['model'].str.contains(model_name)]
    model_dataset = model_dataset[model_dataset['dishonest'] != '[-1]']
    model_dataset.to_csv(dataset_path + model_name + '.csv', index=False)
    print(f"Saved {model_name} dataset to {dataset_path}_{model_name}.csv")
    print(f"Number of statements: {len(model_dataset)}")


dataset_paths = [
    "/home/wombat_share/XXXX-3/dishonesty-bench/results/datasets/convince-me-developer.csv",
    "/home/wombat_share/XXXX-3/dishonesty-bench/results/datasets/harm-pressure-choice.csv",
    "/home/wombat_share/XXXX-3/dishonesty-bench/results/datasets/harm-pressure-knowledge-report.csv",
]

models = [
    # "mistral-small-3.1-24b-instruct",
    # "llama-v3.3-70b-instruct",
    "qwen-2.5-72b-instruct",
]

for dataset_path in dataset_paths:
    for model in models:
        filter_dataset(dataset_path, model)
