JetBrains-Research/commit-message-editing

1import os2 3from datasets import load_dataset4 5CACHE_DIR = 'cache'6N_SAMPLES = 157REMOVED_COMMITS = ['9cc896202dc38d962c01aa2637dbc5bbc3e3dd9b']8 9 10def load_data():11    df = load_dataset("JetBrains-Research/commit-rewriting-samples",12                      split="train",13                      token=os.environ.get('HF_REWRITING_TOKEN'),14                      cache_dir=CACHE_DIR).to_pandas()15 16    removed_idx = df['hash'].isin(REMOVED_COMMITS)17    df = df[~removed_idx]18 19    return df.to_dict('records')[:N_SAMPLES]20