Switch to the special commit rewriting dataset

app.py

@@ -1,3 +1,4 @@
+import json
 import os
 import random
 import uuid
@@ -19,7 +20,9 @@ n_samples = len(data)
 saver = get_dataset_saver(HF_TOKEN, HF_DATASET, private=True)
 
 
-def convert_diff_to_unified(diff):
+def convert_diff_to_unified(diff_string):
+    diff = json.loads(diff_string)
+
     result = "\n".join(
         [
             f'--- {modified_file["old_path"]}\n'

data_loader.py

@@ -1,22 +1,15 @@
+import os
+
 from datasets import load_dataset
 
-MODEL = 'cmg_gpt_4_0613'
 CACHE_DIR = 'cache'
+N_SAMPLES = 15
 
 
 def load_data():
-    dataset = load_dataset("JetBrains-Research/lca-cmg",
-                           "commitchronicle-py-long",
-                           split="test",
-                           cache_dir=CACHE_DIR).to_pandas().set_index(['hash', 'repo']).rename(
-        columns={'message': 'reference'})
-
-    model_dataset = load_dataset("JetBrains-Research/lca-results",
-                                 MODEL,
-                                 split="test",
-                                 cache_dir=CACHE_DIR).to_pandas().set_index(['hash', 'repo'])[["prediction"]]
-
-    model_dataset = model_dataset[~model_dataset.index.duplicated(keep='first')]
-    dataset = dataset.join(other=model_dataset)
+    df = load_dataset("petrtsv-jb/commit-rewriting-samples",
+                      split="train",
+                      token=os.environ.get('HF_REWRITING_TOKEN'),
+                      cache_dir=CACHE_DIR).to_pandas()
 
-    return dataset.reset_index().to_dict('records')
+    return df.to_dict('records')[:N_SAMPLES]