Synthetic dataset generation

change_visualizer.py

@@ -2,28 +2,41 @@ import gradio as gr
 
 import generate_annotated_diffs
 
-df = generate_annotated_diffs.data_with_annotated_diffs()
-n_diffs = len(df)
+df_manual = generate_annotated_diffs.manual_data_with_annotated_diffs()
+n_diffs_manual = len(df_manual)
 
 
-def update_view(diff_idx):
+def update_manual_view(diff_idx):
     diff_idx -= 1
-    return df.iloc[diff_idx]['annotated_diff'], df.iloc[diff_idx]['commit_msg_start'], df.iloc[diff_idx][
-        'commit_msg_end'], df.iloc[diff_idx][
-        'session'], f"https://github.com/{df.iloc[diff_idx]['repo']}/commit/{df.iloc[diff_idx]['hash']}"
+    return df_manual.iloc[diff_idx]['annotated_diff'], df_manual.iloc[diff_idx]['commit_msg_start'], \
+        df_manual.iloc[diff_idx][
+            'commit_msg_end'], df_manual.iloc[diff_idx][
+        'session'], f"https://github.com/{df_manual.iloc[diff_idx]['repo']}/commit/{df_manual.iloc[diff_idx]['hash']}"
 
 
 if __name__ == '__main__':
     with gr.Blocks(theme=gr.themes.Soft()) as application:
-        slider = gr.Slider(minimum=1, maximum=n_diffs, step=1, value=1, label=f"Sample number (total: {n_diffs})")
+        with gr.Tab("Manual"):
+            slider_manual = gr.Slider(minimum=1, maximum=n_diffs_manual, step=1, value=1,
+                                      label=f"Sample number (total: {n_diffs_manual})")
+
+            diff_view_manual = gr.Highlightedtext(combine_adjacent=True, color_map={'+': "green", '-': "red"})
+            start_view_manual = gr.Textbox(interactive=False, label="Start message", container=True)
+            end_view_manual = gr.Textbox(interactive=False, label="End message", container=True)
+            session_view_manual = gr.Textbox(interactive=False, label="Session", container=True)
+            link_view_manual = gr.Markdown()
+            view_manual = [
+                diff_view_manual,
+                start_view_manual,
+                end_view_manual,
+                session_view_manual,
+                link_view_manual
+            ]
+
+            slider_manual.change(update_manual_view, inputs=slider_manual,
+                                 outputs=view_manual)
+
+        application.load(update_manual_view, inputs=slider_manual,
+                         outputs=view_manual)
 
-        diff_view = gr.Highlightedtext(combine_adjacent=True, color_map={'+': "green", '-': "red"})
-        start_view = gr.Textbox(interactive=False, label="Start message", container=True)
-        end_view = gr.Textbox(interactive=False, label="End message", container=True)
-        session_view = gr.Textbox(interactive=False, label="Session", container=True)
-        link_view = gr.Markdown()
-
-        slider.change(update_view, inputs=slider, outputs=[diff_view, start_view, end_view, session_view, link_view])
-
-        application.load(update_view, inputs=slider, outputs=[diff_view, start_view, end_view, session_view, link_view])
     application.launch()

config.py

@@ -1,6 +1,8 @@
 import os
 from pathlib import Path
 
+GRAZIE_API_JWT_TOKEN = os.environ.get("GRAZIE_API_JWT_TOKEN")
+
 HF_TOKEN = os.environ.get('HF_TOKEN')
 HF_RAW_DATASET_NAME = "petrtsv-jb/commit-msg-rewriting"
 HF_RAW_DATASET_SPLIT = 'train'
@@ -11,4 +13,4 @@ CACHE_DIR.mkdir(exist_ok=True)
 OUTPUT_DIR = Path("output")
 OUTPUT_DIR.mkdir(exist_ok=True)
 
-ANNOTATED_DIFFS_ARTIFACT = OUTPUT_DIR / "annotated_diffs.csv"
+SYNTHETIC_DATASET_ARTIFACT = OUTPUT_DIR / "synthetic.csv"

generate_annotated_diffs.py

@@ -1,49 +1,8 @@
-from datetime import datetime
-
 import diff_match_patch as dmp_module
 
 import hf_data_loader
 
 
-def group_changes(changes):
-    groups = {}
-    for change in changes:
-        group = datetime.fromisoformat(change['ts'])
-        if group not in groups:
-            groups[group] = []
-        groups[group].append(change)
-
-    grouped_changes = []
-    for group in sorted(groups.keys()):
-        groups[group].sort(key=lambda x: x['p'])
-        grouped_changes.append(groups[group])
-
-    return grouped_changes
-
-
-def fill_in_annotation_gaps(annotated_text):
-    seg_start = None
-    seg_type = None
-
-    for i, e in enumerate(annotated_text):
-        if e[1] is None:
-            continue
-
-        if seg_type is None:
-            seg_start = i
-        elif seg_type != e[1]:
-            for j in range(seg_start, i):
-                annotated_text[j][1] = seg_type
-            seg_start = i
-            seg_type = e[1]
-
-    if seg_start is not None:
-        for j in range(seg_start, len(annotated_text)):
-            annotated_text[j][1] = seg_type
-
-    return annotated_text
-
-
 def get_annotated_diff(start_text, end_text):
     dmp = dmp_module.diff_match_patch()
     dmp_mapping = {
@@ -60,14 +19,20 @@ def get_annotated_diff(start_text, end_text):
     return result
 
 
-def annotated_diff_for_row(row):
+def annotated_diff_for_row_manual_df(row):
     start = row['commit_msg_start']
     end = row['commit_msg_end']
     return get_annotated_diff(start, end)
 
 
-def data_with_annotated_diffs():
-    df = hf_data_loader.load_raw_dataset_as_pandas()
-    annotated = df.apply(annotated_diff_for_row, axis=1)
+def annotated_diff_for_row_synthetic_df(row):
+    start = row['initial_msg_pred']
+    end = row['reference']
+    return get_annotated_diff(start, end)
+
+
+def manual_data_with_annotated_diffs():
+    df = hf_data_loader.load_raw_rewriting_dataset_as_pandas()
+    annotated = df.apply(annotated_diff_for_row_manual_df, axis=1)
     df['annotated_diff'] = annotated
     return df

generate_synthetic_dataset.py

@@ -0,0 +1,69 @@
+from grazie.api.client.chat.prompt import ChatPrompt
+from grazie.api.client.endpoints import GrazieApiGatewayUrls
+from grazie.api.client.gateway import GrazieApiGatewayClient, GrazieAgent, AuthType
+from grazie.api.client.profiles import LLMProfile
+from tqdm import tqdm
+
+import config
+import hf_data_loader
+
+client = GrazieApiGatewayClient(
+    grazie_agent=GrazieAgent(name="commit-rewriting-summary-generation", version="dev"),
+    url=GrazieApiGatewayUrls.STAGING,
+    auth_type=AuthType.SERVICE,
+    grazie_jwt_token=config.GRAZIE_API_JWT_TOKEN
+)
+
+
+def build_prompt(reference, diff):
+    return f"""A software developer uses a LLM to generate commit messages.
+    
+They generated a commit message for the following source code changes:
+START OF THE SOURCE CODE CHANGES
+{diff}
+END OF THE SOURCE CODE CHANGES
+
+After generating the commit message the developer understands that it is not perfect. After making dome changes,
+they come up with an edited version of the message. Here is this edited message:
+START OF THE COMMIT MESSAGE 
+{reference}
+END OF THE COMMIT MESSAGE
+
+Your task is to print the initial, LLM-generated commit message. Print only the initial commit message's text after the 
+token "OUTPUT".
+
+OUTPUT"""
+
+
+def generate_prompt_for_row(row):
+    reference = row['reference']
+    diff = row['mods']
+    return build_prompt(reference, diff)
+
+
+def generate_initial_msg(prompt):
+    commit_msg = client.chat(
+        chat=ChatPrompt()
+        .add_system("You are a helpful assistant.")
+        .add_user(prompt),
+        profile=LLMProfile("gpt-4-1106-preview")
+    ).content
+
+    return commit_msg
+
+
+def generate_synthetic_dataset():
+    df = hf_data_loader.load_full_commit_dataset_as_pandas()
+    df['initial_msg_prompt'] = df.apply(generate_prompt_for_row, axis=1)
+    initial_messages_pred = []
+
+    for prompt in tqdm(df['initial_msg_prompt']):
+        initial_messages_pred.append(generate_initial_msg(prompt))
+
+    df['initial_msg_pred'] = initial_messages_pred
+
+    df.to_csv(config.SYNTHETIC_DATASET_ARTIFACT)
+
+
+if __name__ == '__main__':
+    generate_synthetic_dataset()

hf_data_loader.py

@@ -3,11 +3,16 @@ from datasets import load_dataset
 import config
 
 
-def load_raw_dataset_as_pandas():
+def load_raw_rewriting_dataset_as_pandas():
     return load_dataset(config.HF_RAW_DATASET_NAME,
                         split=config.HF_RAW_DATASET_SPLIT,
                         token=config.HF_TOKEN,
                         cache_dir=config.CACHE_DIR).to_pandas()
 
 
-load_raw_dataset_as_pandas()
+def load_full_commit_dataset_as_pandas():
+    return load_dataset("JetBrains-Research/lca-commit-message-generation",
+                        "commitchronicle-py-long",
+                        split="test",
+                        cache_dir=config.CACHE_DIR).to_pandas().rename(
+        columns={'message': 'reference'})