When training on several datasets of very different sizes (English vs Swahili, web vs code), sampling proportionally to size starves the small ones. Multilingual models like mT5 and XLM-R use temperature sampling: . T = 1 is proportional, and as T grows the mixture approaches uniform.
Implement:
mixture_weights(sizes, T): sizes maps dataset names to sizes. Return normalised weights as a dict.epochs(sizes, weights, total): how many passes over each dataset a run of total samples makes, weights[name] * total / sizes[name]temperature_for_cap(sizes, total, max_epochs, grid): the smallest T in grid such that no dataset is repeated more than max_epochs times, or None