Kevius/VSI-Bench-modified

1#!/usr/bin/env python32# /// script3# requires-python = ">=3.10"4# dependencies = [5#     "fastparquet",6#     "pandas",7#     "pathlib",8#     "pyarrow",9# ]10# ///11"""12Create parquet files for config subsets of the VSI-Bench dataset.13* debiased: all examples not pruned by Iterative Bias Pruning (aka VSI-Bench-Debiased)14* pruned: all examples pruned by Iterative Bias Pruning15 16> [!NOTE]17> If you do not pass `index=False`, the parquet files will have a `__index_level_0__` column18"""19 20import pandas as pd21from pathlib import Path22 23script_dir = Path(__file__).parent24pruned_ids_path = script_dir / "pruned_ids.txt"25test_jsonl_path = script_dir / "test.jsonl"26pq_debiased_path = script_dir / "test_debiased.parquet"27pq_pruned_path = script_dir / "test_pruned.parquet"28 29print("Creating parquet files...")30 31print(f"Loading pruned ids from '{pruned_ids_path}'...")32with open(pruned_ids_path, "r") as f:33    pruned_ids = f.read().splitlines()34print(f"    -> Loaded {len(pruned_ids)} pruned ids.")35 36print(f"Loading test data from '{test_jsonl_path}'...")37df = pd.read_json(str(test_jsonl_path), lines=True)38print(f"    -> Loaded {len(df)} examples.")39df["pruned"] = df["id"].astype(str).isin(pruned_ids)40print(f"    -> Added pruned column.")41 42# save the debiased and pruned subsets separately to parquet files43df_debiased = df[~df["pruned"]]44df_pruned = df[df["pruned"]]45 46print(f"Saving debiased examples to '{pq_debiased_path}'...")47df_debiased.to_parquet(pq_debiased_path, index=False)48print(f"    -> Saved {len(df_debiased)} debiased examples.")49 50print(f"Saving pruned examples to '{pq_pruned_path}'...")51df_pruned.to_parquet(pq_pruned_path, index=False)52print(f"    -> Saved {len(df_pruned)} pruned examples.")53 54print("Done.") 55