Data Mixture Temperature

The problem, from first principles

A pretraining corpus mixes sources of very different sizes: huge web crawls, medium code collections and small, high-quality book or math sets. Sampling in proportion to size means the small sets are almost never seen, while sampling uniformly would repeat tiny sets thousands of times. A standard compromise is temperature sampling: raise each source's size to the power 1 / T and renormalize. T = 1 reproduces the natural proportions, larger T flattens toward uniform and T < 1 sharpens toward the largest source. It is the same temperature idea used for softmax sampling, applied to dataset sizes.

From theory to code

Implement mixture_weights and epochs_per_source.

Constraints

• mixture_weights(sizes, temperature) returns weights proportional to sizes ** (1 / temperature), normalized to sum to 1. sizes are positive token counts and temperature > 0.
• epochs_per_source(sizes, weights, total_tokens) returns, per source, how many times its data is repeated: weights * total_tokens / sizes.
• Return NumPy float arrays.

Hints