File size: 1,869 Bytes
6676c5a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 |
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()
|