| 196 | |
| 197 | |
| 198 | def get_author_map(data_folder: Path, regen_data, assert_stored=False): |
| 199 | if not regen_data and Path(data_folder / "author_map.pkl").exists(): |
| 200 | with open(data_folder / "author_map.pkl", "rb") as f: |
| 201 | return pickle.load(f) |
| 202 | else: |
| 203 | if assert_stored: |
| 204 | raise FileNotFoundError( |
| 205 | "Author map not found, you are loading for inference you need to have an author map!" |
| 206 | ) |
| 207 | print("Regenerating Author Map") |
| 208 | all_data = pd.read_csv(data_folder / "commitlist.csv") |
| 209 | authors = all_data.author.unique().tolist() |
| 210 | authors.append(UNKNOWN_TOKEN) |
| 211 | author_map = {author: i for i, author in enumerate(authors)} |
| 212 | with open(data_folder / "author_map.pkl", "wb") as f: |
| 213 | pickle.dump(author_map, f) |
| 214 | return author_map |
| 215 | |
| 216 | |
| 217 | def get_file_map(data_folder: Path, regen_data, assert_stored=False): |