Fix the synthetic data generation pipeline

api_wrappers/grazie_wrapper.py

@@ -32,7 +32,7 @@ def llm_request(prompt):
 
     while output is None:
         try:
-            output = output = client.chat(
+            output = client.chat(
                 chat=ChatPrompt()
                 .add_system("You are a helpful assistant.")
                 .add_user(prompt),

dataset_statistics.py

@@ -9,10 +9,7 @@ from scipy.stats import stats
 import config
 
 
-def get_statistics(row):
-    start_msg = row["commit_msg_start"]
-    end_msg = row["commit_msg_end"]
-
+def get_statistics_for_sample(start_msg, end_msg, row=None):
     edit_ops = Levenshtein.editops(start_msg, end_msg)
     n_deletes = sum([1 if op == 'delete' else 0 for op, _, _ in edit_ops])
     n_inserts = sum([1 if op == 'insert' else 0 for op, _, _ in edit_ops])
@@ -32,12 +29,18 @@ def get_statistics(row):
         "changes_norm": n_changes / len(end_msg),
 
         "lendiff": abs(len(start_msg) - len(end_msg)),
-        "editdist": row["editdist_related"]
+        "editdist": row["editdist_related"] if row is not None else Levenshtein.distance(start_msg, end_msg),
     }
 
 
+def get_statistics_for_row(row):
+    start_msg = row["commit_msg_start"]
+    end_msg = row["commit_msg_end"]
+    return get_statistics_for_sample(start_msg, end_msg, row=row)
+
+
 def get_statistics_for_df(df: pd.DataFrame):
-    stats = [get_statistics(row) for _, row in
+    stats = [get_statistics_for_row(row) for _, row in
              df.iterrows()]
 
     assert len(stats) > 0

generation_steps/synthetic_end_to_start.py

@@ -4,8 +4,8 @@ import pandas as pd
 from tqdm import tqdm
 
 import config
-import generate_annotated_diffs
 import dataset_statistics
+import generate_annotated_diffs
 from api_wrappers import grazie_wrapper, hf_data_loader
 from generation_steps import examples
 
@@ -49,9 +49,8 @@ def generate_start_msg(end_msg, diff):
     for i in range(GENERATION_ATTEMPTS):
         start_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 = dataset_statistics.get_statistics_for_sample(start_msg=start_msg_pred, end_msg=end_msg,)
+
         if stats["insertions"] < REL_INSERTIONS_THRESHOLD:
             return start_msg_pred
         else:

generation_steps/synthetic_start_to_end.py

@@ -2,7 +2,6 @@ import pandas as pd
 from tqdm import tqdm
 
 import config
-import generate_annotated_diffs
 import dataset_statistics
 from api_wrappers import grazie_wrapper
 from generation_steps import examples
@@ -47,9 +46,7 @@ def generate_end_msg(start_msg, diff):
     for i in range(GENERATION_ATTEMPTS):
         end_msg_pred = grazie_wrapper.generate_for_prompt(prompt)
 
-        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))
+        stats = dataset_statistics.get_statistics_for_sample(start_msg=start_msg, end_msg=end_msg_pred, )
         if stats["deletions"] < REL_DELETIONS_THRESHOLD:
             return end_msg_pred
         else: