| 302 | |
| 303 | |
| 304 | def gen_class_weights(dataset: List): |
| 305 | from collections import Counter |
| 306 | |
| 307 | epsilon = 1e-1 |
| 308 | title, files, author, category = zip(*dataset) |
| 309 | category = [common.categories.index(cat) for cat in category] |
| 310 | counter = Counter(category) |
| 311 | percentile_33 = len(category) // 3 |
| 312 | most_common = counter.most_common(percentile_33) |
| 313 | least_common = counter.most_common()[-percentile_33:] |
| 314 | smoothed_top = sum(i[1] + epsilon for i in most_common) / len(most_common) |
| 315 | smoothed_bottom = sum(i[1] + epsilon for i in least_common) / len(least_common) // 3 |
| 316 | class_weights = torch.tensor( |
| 317 | [ |
| 318 | 1.0 / (min(max(counter[i], smoothed_bottom), smoothed_top) + epsilon) |
| 319 | for i in range(len(common.categories)) |
| 320 | ], |
| 321 | device=device, |
| 322 | ) |
| 323 | return class_weights |
| 324 | |
| 325 | |
| 326 | def train(save_path: Path, data_folder: Path, regen_data: bool, resample: bool): |