Generate a dataset for the labeling app

api_wrappers/grazie_wrapper.py

@@ -10,7 +10,7 @@ import config
 client = GrazieApiGatewayClient(
     grazie_agent=GrazieAgent(name="commit-rewriting-synthetic-end-to-start", version="dev"),
     url=GrazieApiGatewayUrls.STAGING,
-    auth_type=AuthType.SERVICE,
+    auth_type=AuthType.USER,
     grazie_jwt_token=config.GRAZIE_API_JWT_TOKEN
 )
 

api_wrappers/hf_data_loader.py

@@ -3,14 +3,14 @@ from datasets import load_dataset
 import config
 
 
-def load_raw_rewriting_dataset_as_pandas():
+def load_raw_rewriting_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()
 
 
-def load_full_commit_dataset_as_pandas():
+def load_full_commit_as_pandas():
     return load_dataset(path=config.HF_FULL_COMMITS_DATASET_NAME,
                         name=config.HF_FULL_COMMITS_DATASET_SUBNAME,
                         split=config.HF_FULL_COMMITS_DATASET_SPLIT,
@@ -18,19 +18,35 @@ def load_full_commit_dataset_as_pandas():
         columns={'message': 'reference'})
 
 
-def load_processed_rewriting_dataset_as_pandas():
-    manual_rewriting = load_raw_rewriting_dataset_as_pandas()[
+def load_processed_rewriting_as_pandas():
+    manual_rewriting = load_raw_rewriting_as_pandas()[
         ["hash", "repo", "commit_msg_start", "commit_msg_end", "session"]]
     manual_rewriting.set_index(["hash", "repo"], inplace=True)
 
-    mods_dataset = load_full_commit_dataset_as_pandas()[["hash", "repo", "mods"]]
+    mods_dataset = load_full_commit_as_pandas()[["hash", "repo", "mods"]]
     mods_dataset.set_index(["hash", "repo"], inplace=True)
 
     return manual_rewriting.join(other=mods_dataset, how='left').reset_index()
 
 
-def load_synthetic_dataset_as_pandas():
+def load_synthetic_as_pandas():
     return load_dataset(config.HF_SYNTHETIC_DATASET_NAME,
                         split=config.HF_SYNTHETIC_DATASET_SPLIT,
                         token=config.HF_TOKEN,
                         cache_dir=config.CACHE_DIR).to_pandas()
+
+
+def load_full_commit_with_predictions_as_pandas():
+    full_dataset = load_full_commit_as_pandas()
+    predictions_dataset = load_dataset(config.HF_PREDICTIONS_DATASET_NAME,
+                                       config.HF_PREDICTIONS_DATASET_SUBNAME,
+                                       split=config.HF_PREDICTIONS_DATASET_SPLIT,
+                                       cache_dir=config.CACHE_DIR
+                                       ).to_pandas().sample(frac=1, random_state=config.RANDOM_STATE
+                                                            ).set_index(['hash', 'repo'])[["prediction"]]
+
+    predictions_dataset = predictions_dataset[~predictions_dataset.index.duplicated(keep='first')]
+
+    dataset = full_dataset.join(other=predictions_dataset, on=('hash', 'repo'))
+
+    return dataset.reset_index()

config.py

@@ -15,6 +15,10 @@ HF_FULL_COMMITS_DATASET_NAME = "JetBrains-Research/lca-commit-message-generation
 HF_FULL_COMMITS_DATASET_SUBNAME = "commitchronicle-py-long"
 HF_FULL_COMMITS_DATASET_SPLIT = "test"
 
+HF_PREDICTIONS_DATASET_NAME = "JetBrains-Research/lca-results"
+HF_PREDICTIONS_DATASET_SUBNAME = "cmg_gpt_4_0613"
+HF_PREDICTIONS_DATASET_SPLIT = "test"
+
 HF_SYNTHETIC_DATASET_NAME = "petrtsv-jb/synthetic-commit-msg-rewriting"
 HF_SYNTHETIC_DATASET_SPLIT = 'train'
 
@@ -24,8 +28,8 @@ CACHE_DIR.mkdir(exist_ok=True)
 OUTPUT_DIR = Path("output")
 OUTPUT_DIR.mkdir(exist_ok=True)
 
-
 END_TO_START_ARTIFACT = OUTPUT_DIR / "end_to_start.csv"
 START_TO_END_ARTIFACT = OUTPUT_DIR / "start_to_end.csv"
 SYNTHETIC_DATASET_ARTIFACT = OUTPUT_DIR / "synthetic.csv"
 METRICS_CORRELATIONS_ARTIFACT = OUTPUT_DIR / "metrics_correlations.csv"
+DATA_FOR_LABELING_ARTIFACT = OUTPUT_DIR / "data_for_labeling.csv"

generate_annotated_diffs.py

@@ -26,14 +26,14 @@ def annotated_diff_for_row(row):
 
 
 def manual_data_with_annotated_diffs():
-    df = hf_data_loader.load_raw_rewriting_dataset_as_pandas()
+    df = hf_data_loader.load_raw_rewriting_as_pandas()
     annotated = df.apply(annotated_diff_for_row, axis=1)
     df['annotated_diff'] = annotated
     return df
 
 
 def synthetic_data_with_annotated_diffs():
-    df = hf_data_loader.load_synthetic_dataset_as_pandas()
+    df = hf_data_loader.load_synthetic_as_pandas()
     annotated = df.apply(annotated_diff_for_row, axis=1)
     df['annotated_diff'] = annotated
     return df

generate_synthetic_dataset.py

@@ -4,7 +4,7 @@ from generation_steps import synthetic_end_to_start, synthetic_start_to_end, met
 
 
 def run():
-    df = hf_data_loader.load_processed_rewriting_dataset_as_pandas()
+    df = hf_data_loader.load_processed_rewriting_as_pandas()
 
     df = synthetic_end_to_start.transform(df)
     df = synthetic_start_to_end.transform(df)

generation_steps/examples.py

@@ -36,7 +36,7 @@ END OF THE IMPROVED COMMIT MESSAGE
 END OF THE EXAMPLE"""
 
 
-manual_df = hf_data_loader.load_raw_rewriting_dataset_as_pandas()[['commit_msg_start', 'commit_msg_end']]
+manual_df = hf_data_loader.load_raw_rewriting_as_pandas()[['commit_msg_start', 'commit_msg_end']]
 manual_df = manual_df.sample(n=N_EXAMPLES, random_state=config.RANDOM_STATE)
 
 

generation_steps/for_labeling.py

@@ -0,0 +1,58 @@
+import json
+
+from tqdm import tqdm
+
+import config
+from api_wrappers import hf_data_loader
+from generation_steps import synthetic_start_to_end
+
+
+def transform(df):
+    print(f"Generating data for labeling:")
+    synthetic_start_to_end.print_config()
+    tqdm.pandas()
+
+    manual_df = hf_data_loader.load_raw_rewriting_as_pandas()
+
+    manual_df = manual_df.sample(frac=1, random_state=config.RANDOM_STATE
+                                 ).set_index(['hash', 'repo'])[["commit_msg_start", "commit_msg_end"]]
+
+    manual_df = manual_df[~manual_df.index.duplicated(keep='first')]
+
+    def get_is_manually_rewritten(row):
+        commit_id = (row['hash'], row['repo'])
+        return commit_id in manual_df.index
+
+    result = df
+    result['manual_sample'] = result.progress_apply(get_is_manually_rewritten, axis=1)
+
+    def get_prediction_message(row):
+        commit_id = (row['hash'], row['repo'])
+        if row['manual_sample']:
+            return manual_df.loc[commit_id]['commit_msg_start']
+        return row['prediction']
+
+    def get_enhanced_message(row):
+        commit_id = (row['hash'], row['repo'])
+        if row['manual_sample']:
+            return manual_df.loc[commit_id]['commit_msg_end']
+        return synthetic_start_to_end.generate_end_msg(start_msg=row["prediction"],
+                                                       diff=row["mods"])
+
+    result['enhanced'] = result.progress_apply(get_enhanced_message, axis=1)
+    result['prediction'] = result.progress_apply(get_prediction_message, axis=1)
+    result['mods'] = result['mods'].progress_apply(json.dumps)
+
+    result.to_csv(config.DATA_FOR_LABELING_ARTIFACT)
+    print("Done")
+    return result
+
+
+def main():
+    synthetic_start_to_end.GENERATION_ATTEMPTS = 3
+    df = hf_data_loader.load_full_commit_with_predictions_as_pandas()
+    transform(df)
+
+
+if __name__ == '__main__':
+    main()

generation_steps/metrics_analysis.py

@@ -77,7 +77,7 @@ METRICS = {
 
 
 def attach_references(df):
-    reference_df = hf_data_loader.load_full_commit_dataset_as_pandas().set_index(["hash", "repo"])[["reference"]]
+    reference_df = hf_data_loader.load_full_commit_as_pandas().set_index(["hash", "repo"])[["reference"]]
     df = df.set_index(["hash", "repo"])
     return df.join(other=reference_df, how="left").reset_index()
 

generation_steps/synthetic_end_to_start.py

@@ -98,7 +98,7 @@ def transform(df):
 
 
 def main():
-    df = hf_data_loader.load_processed_rewriting_dataset_as_pandas()
+    df = hf_data_loader.load_processed_rewriting_as_pandas()
     transform(df)
 
 

generation_steps/synthetic_start_to_end.py

@@ -12,7 +12,7 @@ REL_DELETIONS_THRESHOLD = 0.75
 GENERATION_ATTEMPTS = 5
 
 
-def build_prompt(reference, diff):
+def build_prompt(prediction, diff):
     return f"""A LLM generated a commit message for the following source code changes:
 START OF THE SOURCE CODE CHANGES
 {diff}
@@ -20,7 +20,7 @@ END OF THE SOURCE CODE CHANGES
 
 Here is the message the LLM generated:
 START OF THE COMMIT MESSAGE 
-{reference}
+{prediction}
 END OF THE COMMIT MESSAGE
 
 This generated message is not perfect. Your task is to rewrite and improve it.
@@ -40,20 +40,20 @@ token "OUTPUT".
 OUTPUT"""
 
 
-def generate_start_msg(end_msg, diff):
-    prompt = build_prompt(reference=end_msg, diff=diff)
+def generate_end_msg(start_msg, diff):
+    prompt = build_prompt(prediction=start_msg, diff=diff)
     results = []
 
     for i in range(GENERATION_ATTEMPTS):
-        start_msg_pred = grazie_wrapper.generate_for_prompt(prompt)
+        end_msg_pred = grazie_wrapper.generate_for_prompt(prompt)
 
-        stats = statistics.get_statistics(start_msg=start_msg_pred, end_msg=end_msg,
-                                          annotated_msg=generate_annotated_diffs.get_annotated_diff(start_msg_pred,
-                                                                                                    end_msg))
+        stats = statistics.get_statistics(start_msg=start_msg, end_msg=end_msg_pred,
+                                          annotated_msg=generate_annotated_diffs.get_annotated_diff(start_msg,
+                                                                                                    end_msg_pred))
         if stats["deletions"] < REL_DELETIONS_THRESHOLD:
-            return start_msg_pred
+            return end_msg_pred
         else:
-            results.append((stats["deletions"], start_msg_pred))
+            results.append((stats["deletions"], end_msg_pred))
 
     results.sort()
     return results[0][1]
@@ -62,13 +62,17 @@ def generate_start_msg(end_msg, diff):
 COLS_TO_KEEP = ["hash", "repo", "commit_msg_start", "mods", "session", "end_to_start"]
 
 
-def transform(df):
-    print(f"Start -> send synthesis:")
+def print_config():
     print(f"NUMBER OF EXAMPLES PER PROMPT = {examples.N_EXAMPLES}")
     print(f"GENERATION_MULTIPLIER = {GENERATION_MULTIPLIER}")
     print(f"REL_DELETIONS_THRESHOLD = {REL_DELETIONS_THRESHOLD}")
     print(f"GENERATION_ATTEMPTS = {GENERATION_ATTEMPTS}")
 
+
+def transform(df):
+    print(f"Start -> send synthesis:")
+    print_config()
+
     df['start_to_end'] = False
 
     generated_data = {
@@ -80,8 +84,8 @@ def transform(df):
 
     for _, row in tqdm(df.iterrows(), total=len(df)):
         for i in range(GENERATION_MULTIPLIER):
-            commit_msg_end_pred = generate_start_msg(end_msg=row["commit_msg_start"],
-                                                     diff=row["mods"])
+            commit_msg_end_pred = generate_end_msg(start_msg=row["commit_msg_start"],
+                                                   diff=row["mods"])
 
             generated_data["commit_msg_end"].append(commit_msg_end_pred)
             for col in COLS_TO_KEEP: