Repro: SwissDecisionSummaryTranslations test split is 1975-1982 due to constant-key unstable sort (JoelNiklaus/SwissLegalTranslations#1)

swiltra_split_repro.py

"""Reproduce the SwissDecisionSummaryTranslations val/test split bug.

The released headnote test split (2,000 BGEs) is almost entirely BGE volumes
101-108 (1975-1982) instead of the recency-based split described in the
SwiLTra-Bench paper. Cause: at the code state that generated the data
(JoelNiklaus/SwissLegalTranslations @ 761fc5c8, 2024-11-21), get_split_cols
sorted twice:

    df = df.sort_values(by='year', ascending=False)        # intended recency
    ...
    df = df.sort_values(by='num_langs', ascending=False)   # constant key!

Every BGE has all three languages, so num_langs == 3 everywhere. pandas'
default quicksort is not stable, so the second sort destroys the year
ordering before create_split_configs slices the first 2,000 unique ids.
Fixed on main in ff823006 (2024-11-25), but the headnote parquets were last
generated 2024-11-21 and never rebuilt.

This script downloads the released parquets, rebuilds the 19,001-BGE
universe in the order the pipeline saw it (df.pivot emits bge-ascending),
applies the buggy two-sort logic, and compares the resulting first-2,000
slice with the published test split. Expected output (pandas 2.2.3,
numpy introsort):

    buggy two-sort (as shipped):   overlap 2000/2000   test years 101-148 mean 104.7
    current main (combined sort):  overlap    1/2000   test years 139-148 mean 143.1
    stable second sort (control):  overlap    1/2000

Requires: pandas, pyarrow. Verified with pandas 2.2.3 / pyarrow 18.1.0.
"""

import urllib.request

import pandas as pd

BASE = (
    "https://huggingface.co/datasets/joelniklaus/"
    "SwissDecisionSummaryTranslations/resolve/main/bge_level"
)

def fetch(split: str) -> pd.DataFrame:
    fname = f"bge_{split}.parquet"
    urllib.request.urlretrieve(f"{BASE}/{split}.parquet", fname)
    return pd.read_parquet(fname)

train, val, test = fetch("train"), fetch("val"), fetch("test")
universe = pd.concat([train, val, test], ignore_index=True)
print(f"universe: {len(universe)} BGEs, volumes {universe.year.min()}-{universe.year.max()}")
print(f"published test: n={len(test)}, volumes {test.year.min()}-{test.year.max()}, "
      f"mean {test.year.mean():.1f}\n")

lang_cols = [c for c in universe.columns if c.endswith("_bgeText")]
TEST_SIZE = 2000

def first_2000(df: pd.DataFrame, combined: bool, second_kind: str = "quicksort") -> set:
    d = df.copy()
    d["num_langs"] = sum((d[c] != "").astype(int) for c in lang_cols)
    # df.pivot(index=group_cols, ...) in combine_rows emits rows sorted by the
    # index, i.e. bge-ascending; reconstruct that input order.
    d = d.sort_values("bge", kind="stable").reset_index(drop=True)
    if combined:
        # utils.py @ ff823006 (current main)
        d = d.sort_values(by=["num_langs", "year"], ascending=[False, False])
    else:
        # utils.py @ 761fc5c8 (generated the released data)
        d = d.sort_values(by="year", ascending=False)
        d = d.sort_values(by="num_langs", ascending=False, kind=second_kind)
    # create_split_configs: unique ids, first TEST_SIZE -> test
    return set(list(dict.fromkeys(d.bge))[:TEST_SIZE])

published = set(test.bge)
for label, kwargs in [
    ("buggy two-sort (as shipped)  ", dict(combined=False)),
    ("current main (combined sort) ", dict(combined=True)),
    ("stable second sort (control) ", dict(combined=False, second_kind="stable")),
]:
    pred = first_2000(universe, **kwargs)
    years = universe[universe.bge.isin(pred)].year
    print(f"{label} overlap {len(pred & published):4d}/2000   "
          f"pred-test years {years.min()}-{years.max()} mean {years.mean():.1f}")