Keep the session column

api_wrappers/hf_data_loader.py

@@ -19,7 +19,8 @@ def load_full_commit_dataset_as_pandas():
 
 
 def load_processed_rewriting_dataset_as_pandas():
-    manual_rewriting = load_raw_rewriting_dataset_as_pandas()[["hash", "repo", "commit_msg_start", "commit_msg_end"]]
+    manual_rewriting = load_raw_rewriting_dataset_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"]]

change_visualizer.py

@@ -15,10 +15,11 @@ STATISTICS = {"manual": statistics.get_statistics_for_df(df_manual),
 
 def update_dataset_view(diff_idx, df):
     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.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']}")
 
 
 def update_dataset_view_manual(diff_idx):

generation_steps/synthetic_end_to_start.py

@@ -3,8 +3,8 @@ from tqdm import tqdm
 
 import config
 import generate_annotated_diffs
-from api_wrappers import grazie_wrapper, hf_data_loader
 import statistics
+from api_wrappers import grazie_wrapper, hf_data_loader
 
 N_EXAMPLES = 5
 GENERATION_MULTIPLIER = 2
@@ -89,26 +89,27 @@ def generate_start_msg(end_msg, diff):
     return results[0][1]
 
 
+COLS_TO_KEEP = ["hash", "repo", "commit_msg_end", "mods", "session"]
+
+
 def transform(df):
     df['end_to_start'] = False
 
     generated_data = {
-        "hash": [],
-        "repo": [],
-        "commit_msg_start": [],
-        "commit_msg_end": [],
-        "mods": []
+        "commit_msg_start": []
     }
 
+    for col in COLS_TO_KEEP:
+        generated_data[col] = []
+
     for _, row in tqdm(df.iterrows(), total=len(df)):
         for i in range(GENERATION_MULTIPLIER):
             commit_msg_start_pred = generate_start_msg(end_msg=row["commit_msg_end"],
                                                        diff=row["mods"])
-            generated_data["hash"].append(row["hash"])
-            generated_data["repo"].append(row["repo"])
+
             generated_data["commit_msg_start"].append(commit_msg_start_pred)
-            generated_data["commit_msg_end"].append(row["commit_msg_end"])
-            generated_data["mods"].append(row["mods"])
+            for col in COLS_TO_KEEP:
+                generated_data[col].append(row[col])
 
     generated_df = pd.DataFrame.from_dict(generated_data)
     generated_df['end_to_start'] = True