JetBrains-Research/commit-message-editing-visualization

1import config2from api_wrappers import hf_data_loader3 4N_EXAMPLES = 155 6 7def get_example_prompt_end_to_start(start_msg, end_msg):8    return f"""START OF THE EXAMPLE9 10For the following edited commit message:11START OF THE EDITED COMMIT MESSAGE 12{end_msg}13END OF THE EDITED COMMIT MESSAGE14 15You would output the following initial commit message:16START OF THE INITIAL COMMIT MESSAGE17{start_msg}18END OF THE INITIAL COMMIT MESSAGE19 20END OF THE EXAMPLE"""21 22 23def get_example_prompt_start_to_end(start_msg, end_msg):24    return f"""START OF THE EXAMPLE25 26For the following LLM-generated commit message:27START OF THE GENERATED COMMIT MESSAGE 28{start_msg}29END OF THE GENERATED COMMIT MESSAGE30 31You would output the following improved commit message:32START OF THE IMPROVED COMMIT MESSAGE33{end_msg}34END OF THE IMPROVED COMMIT MESSAGE35 36END OF THE EXAMPLE"""37 38 39manual_df = hf_data_loader.load_raw_rewriting_as_pandas()[["commit_msg_start", "commit_msg_end"]]40manual_df = manual_df.sample(n=N_EXAMPLES, random_state=config.RANDOM_STATE)41 42 43def generate_examples(end_to_start):44    prompt_fn = get_example_prompt_end_to_start if end_to_start else get_example_prompt_start_to_end45    examples = [prompt_fn(row["commit_msg_start"], row["commit_msg_end"]) for _, row in manual_df.iterrows()]46 47    return "\n".join(examples)48 49 50EXAMPLES_END_TO_START = generate_examples(end_to_start=True)51EXAMPLES_START_TO_END = generate_examples(end_to_start=False)52