generate_annotated_diffs.py

import diff_match_patch as dmp_module
from tqdm import tqdm

from api_wrappers import hf_data_loader


def get_annotated_diff(start_text, end_text):
    dmp = dmp_module.diff_match_patch()
    dmp_mapping = {
        -1: '-',
        0: None,
        1: '+'
    }

    diff = dmp.diff_main(start_text, end_text)
    dmp.diff_cleanupSemantic(diff)

    result = [[w, dmp_mapping[t]] for t, w in diff]

    return result


def annotated_diff_for_row(row):
    if "commit_msg_start" in row:
        start = row['commit_msg_start']
    else:
        start = row["G_text"]
    if "commit_msg_end" in row:
        end = row['commit_msg_end']
    else:
        end = row["E_text"]
    return get_annotated_diff(start, end)


def data_with_annotated_diffs():
    tqdm.pandas()

    df = hf_data_loader.load_synthetic_as_pandas()
    df = df.loc[df.is_related].copy()
    annotated = df.progress_apply(annotated_diff_for_row, axis=1)
    df['annotated_diff'] = annotated
    return df