Go through the text data by order of src length with a bit of randomness. From fastai repo.
(data: List, bs: int, shuffle=True)
| 344 | |
| 345 | |
| 346 | def sortish_sampler_indices(data: List, bs: int, shuffle=True) -> np.array: |
| 347 | "Go through the text data by order of src length with a bit of randomness. From fastai repo." |
| 348 | if not shuffle: |
| 349 | return np.argsort(np.array(data) * -1) |
| 350 | |
| 351 | def key_fn(i): |
| 352 | return data[i] |
| 353 | |
| 354 | idxs = np.random.permutation(len(data)) |
| 355 | sz = bs * 50 |
| 356 | ck_idx = [idxs[i: i + sz] for i in range(0, len(idxs), sz)] |
| 357 | sort_idx = np.concatenate([sorted(s, key=key_fn, reverse=True) for s in ck_idx]) |
| 358 | sz = bs |
| 359 | ck_idx = [sort_idx[i: i + sz] for i in range(0, len(sort_idx), sz)] |
| 360 | max_ck = np.argmax([key_fn(ck[0]) for ck in ck_idx]) # find the chunk with the largest key, |
| 361 | ck_idx[0], ck_idx[max_ck] = ck_idx[max_ck], ck_idx[0] # then make sure it goes first. |
| 362 | sort_idx = np.concatenate(np.random.permutation(ck_idx[1:])) if len(ck_idx) > 1 else np.array([], dtype=np.int) |
| 363 | sort_idx = np.concatenate((ck_idx[0], sort_idx)) |
| 364 | return sort_idx |
| 365 | |
| 366 | |
| 367 | class DistributedSortishSampler(Sampler): |