| 164 | |
| 165 | |
| 166 | def get_train_val_data(data_folder: Path, regen_data: bool, train_percentage=0.95): |
| 167 | if ( |
| 168 | not regen_data |
| 169 | and Path(data_folder / "train_df.csv").exists() |
| 170 | and Path(data_folder / "val_df.csv").exists() |
| 171 | ): |
| 172 | train_data = pd.read_csv(data_folder / "train_df.csv") |
| 173 | val_data = pd.read_csv(data_folder / "val_df.csv") |
| 174 | return train_data, val_data |
| 175 | else: |
| 176 | print("Train, Val, Test Split not found generating from scratch.") |
| 177 | commit_list_df = pd.read_csv(data_folder / "commitlist.csv") |
| 178 | test_df = commit_list_df[commit_list_df["category"] == "Uncategorized"] |
| 179 | all_train_df = commit_list_df[commit_list_df["category"] != "Uncategorized"] |
| 180 | # We are going to drop skip from training set since it is so imbalanced |
| 181 | print( |
| 182 | "We are removing skip categories, YOU MIGHT WANT TO CHANGE THIS, BUT THIS IS A MORE HELPFUL CLASSIFIER FOR LABELING." |
| 183 | ) |
| 184 | all_train_df = all_train_df[all_train_df["category"] != "skip"] |
| 185 | all_train_df = all_train_df.sample(frac=1).reset_index(drop=True) |
| 186 | split_index = math.floor(train_percentage * len(all_train_df)) |
| 187 | train_df = all_train_df[:split_index] |
| 188 | val_df = all_train_df[split_index:] |
| 189 | print("Train data size: ", len(train_df)) |
| 190 | print("Val data size: ", len(val_df)) |
| 191 | |
| 192 | test_df.to_csv(data_folder / "test_df.csv", index=False) |
| 193 | train_df.to_csv(data_folder / "train_df.csv", index=False) |
| 194 | val_df.to_csv(data_folder / "val_df.csv", index=False) |
| 195 | return train_df, val_df |
| 196 | |
| 197 | |
| 198 | def get_author_map(data_folder: Path, regen_data, assert_stored=False): |