| 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() | |