JetBrains-Research/commit-message-editing-visualization

1import json2import os3from datetime import datetime, timedelta4 5import pandas as pd6from datasets import load_dataset7from huggingface_hub import hf_hub_download, list_repo_tree8 9import config10 11 12def load_raw_rewriting_as_pandas():13    return load_dataset(14        config.HF_RAW_DATASET_NAME, split=config.HF_RAW_DATASET_SPLIT, token=config.HF_TOKEN, cache_dir=config.CACHE_DIR15    ).to_pandas()16 17 18def load_full_commit_as_pandas():19    return (20        load_dataset(21            path=config.HF_FULL_COMMITS_DATASET_NAME,22            name=config.HF_FULL_COMMITS_DATASET_SUBNAME,23            split=config.HF_FULL_COMMITS_DATASET_SPLIT,24            cache_dir=config.CACHE_DIR,25        )26        .to_pandas()27        .rename(columns={"message": "reference"})28    )29 30 31def edit_time_from_history(history_str):32    history = json.loads(history_str)33 34    if len(history) == 0:35        return 036 37    timestamps = list(map(lambda e: datetime.fromisoformat(e["ts"]), history))38    delta = max(timestamps) - min(timestamps)39 40    return delta // timedelta(milliseconds=1)41 42 43def edit_time_from_timestamps(row):44    loaded_ts = datetime.fromisoformat(row["loaded_ts"])45    submitted_ts = datetime.fromisoformat(row["submitted_ts"])46 47    delta = submitted_ts - loaded_ts48 49    result = delta // timedelta(milliseconds=1)50 51    return result if result >= 0 else None52 53 54def load_processed_rewriting_as_pandas():55    manual_rewriting = load_raw_rewriting_as_pandas()[56        [57            "hash",58            "repo",59            "commit_msg_start",60            "commit_msg_end",61            "session",62            "commit_msg_history",63            "loaded_ts",64            "submitted_ts",65        ]66    ]67 68    manual_rewriting["edit_time_hist"] = manual_rewriting["commit_msg_history"].apply(edit_time_from_history)69    manual_rewriting["edit_time"] = manual_rewriting.apply(edit_time_from_timestamps, axis=1)70 71    manual_rewriting.drop(columns=["commit_msg_history", "loaded_ts", "submitted_ts"])72 73    manual_rewriting.set_index(["hash", "repo"], inplace=True)74 75    mods_dataset = load_full_commit_as_pandas()[["hash", "repo", "mods"]]76    mods_dataset.set_index(["hash", "repo"], inplace=True)77 78    return manual_rewriting.join(other=mods_dataset, how="left").reset_index()79 80 81def load_synthetic_as_pandas():82    return load_dataset(83        config.HF_SYNTHETIC_DATASET_NAME,84        "all_pairs_with_metrics",85        split=config.HF_SYNTHETIC_DATASET_SPLIT,86        token=config.HF_TOKEN,87        cache_dir=config.CACHE_DIR,88    ).to_pandas()89 90 91def load_full_commit_with_predictions_as_pandas():92    full_dataset = load_full_commit_as_pandas()93 94    predictions_paths = []95    for prediction_file in list_repo_tree(96        repo_id=config.HF_PREDICTIONS_DATASET_NAME,97        path=os.path.join("commit_message_generation/predictions", config.HF_PREDICTIONS_MODEL),98        repo_type="dataset",99    ):100        predictions_paths.append(101            hf_hub_download(102                prediction_file.path,103                repo_id=config.HF_PREDICTIONS_DATASET_NAME,104                repo_type="dataset",105                cache_dir=config.CACHE_DIR,106            )107        )108 109    dfs = []110    for path in predictions_paths:111        dfs.append(pd.read_json(path, orient="records", lines=True))112    predictions_dataset = pd.concat(dfs, axis=0, ignore_index=True)113    predictions_dataset = predictions_dataset.sample(frac=1, random_state=config.RANDOM_STATE).set_index(114        ["hash", "repo"]115    )[["prediction"]]116    predictions_dataset = predictions_dataset[~predictions_dataset.index.duplicated(keep="first")]117 118    dataset = full_dataset.join(other=predictions_dataset, on=("hash", "repo"))119 120    return dataset.reset_index()121