Turn raw web text into pretraining data with heuristic filters, global deduplication, eval decontamination, and a quality classifier distilled from LLM labels.
Refine a pretraining corpus
Code available on GitHub.
Most of the work that goes into a language model happens before training starts, on the data. The data teams at LLM labs take raw web crawl, tens of billions of pages, and decide which of them the model will read. Public pipelines show the shape of this work. FineWeb extracts text from Common Crawl, filters it with heuristic rules, and removes near-duplicates. Its educational subset, FineWeb-Edu, then had Llama-3-70B-Instruct score about 450,000 pages for educational value, trained a small classifier on those scores, and used the classifier to filter 15 trillion tokens down to 1.3 trillion. Essential-Web labels 24 trillion tokens with a taxonomy the same way, with a 0.5B model distilled from a 32B one. Zyda-2 deduplicates across source datasets and filters by quality, and OLMo 3 decontaminates its training data against its evaluation suite.
The pattern is the same everywhere: cheap rules on every document, a global step that needs to see every shard at once, an expensive LLM on a small sample, and a cheap model that carries the LLM’s judgment to the rest. At scale it is hard because each stage wants different hardware, a full pass is expensive, and the thresholds change: you want to rerun the last step without paying for the first ones again.
This tutorial runs that pipeline, scaled down to 300,000 FineWeb documents:
- heuristic filters, exact hashes, and MinHash signatures on a warm pool of CPU workers, one cached task per shard;
- a global step that removes exact and near duplicates across shards, and documents that overlap the GSM8K, MMLU, or ARC test sets;
- an open LLM, Qwen2.5-14B-Instruct under vLLM on a warm L40S, that scores 4,000 sampled documents on the FineWeb-Edu rubric;
- a classifier distilled from those labels, trained and run on L4 GPUs, that scores every surviving document;
- measurements of how well the classifier agrees with the LLM and with the public FineWeb-Edu classifier, and what each costs per document;
- the curated corpus, published as an artifact with a dataset card, and a report.
The pipeline
flowchart LR
S["N x refine_shard<br/>warm CPU pool, cached"] --> G["global_dedup<br/>exact, MinHash LSH, eval overlap"]
G --> P["sample_for_labeling"]
P --> L["K x label_chunk<br/>vLLM on a warm L40S"]
L --> T["train_classifier<br/>frozen encoder + head, L4"]
T --> C["N x score_shard<br/>warm L4 pool"]
G --> C
C --> U["curate<br/>threshold"] --> A["artifact:<br/>edu-web-corpus"]
classDef gpu fill:#3b82f633,stroke:#3b82f6,stroke-width:1.5px
classDef cpu fill:#94a3b826,stroke:#94a3b8,stroke-width:1px
class L,T,C gpu
class S,G,P,U,A cpu
| Need | Construct |
|---|---|
| Filter hundreds of shards in parallel without a pod per shard | A reusable CPU environment, flyte.ReusePolicy(replicas=(2, 8)), and asyncio.gather over the shards |
| Change a threshold without redoing the expensive stages | cache="auto" on every stage; the filter settings are a task input, and a version string salts the cache of the stages that use the labeler and classifier modules |
| Keep an LLM loaded between labeling calls | A reusable GPU environment with one L40S replica that keeps the vLLM server running |
| Run a small encoder on a cheaper GPU | A second GPU environment on L4s |
| Pass shards between stages | flyte.io.File for Parquet shards, flyte.io.DataFrame for the result |
| Publish the result with its provenance | flyte.artifacts with a dataset card |
| See what each stage did | A report with the funnel, rejected examples, and the classifier’s agreement |
Getting started
You need a Union cluster with L40S and L4 GPUs. The data and models are public, so no Hugging Face token is needed.
git clone https://github.com/unionai/unionai-examples
cd unionai-examples/v2/tutorials/pretraining_data_refinery
python refinery.py # 15 shards of 20,000 documents, 4,000 LLM labels
python refinery.py --smoke # 2 shards of 2,000 documents, 300 labels
python refinery.py --threshold 2.5 # same corpus, new thresholdImages and environments
The CPU stages need Arrow, xxhash for stable token hashes, and SciPy for connected components. The GPU stages share one image with vLLM, which also brings PyTorch and Transformers for the classifier.
# CPU stages: Arrow for the shards, xxhash for stable token hashes, SciPy for connected components.
cpu_image = flyte.Image.from_debian_base(name="refinery-cpu").with_pip_packages(
"pyarrow==22.0.0", "pandas==2.3.3", "numpy==2.3.5", "scipy==1.16.3", "xxhash==3.6.0",
"fsspec==2025.10.0", "aiohttp==3.13.2", "requests", "unionai-reuse==0.2.1",
)
# GPU stages: vLLM for the labeler. It brings torch and transformers, which the classifier uses.
gpu_image = flyte.Image.from_debian_base(name="refinery-gpu", python_version=(3, 12)).with_pip_packages(
"vllm==0.33.0", "httpx>=0.27", "pandas", "pyarrow", "unionai-reuse==0.2.1",
)
Each stage gets the hardware it needs. The labeler holds one L40S and keeps it: the vLLM server starts on the first call and every later call finds it warm. The classifier is a 110M-parameter encoder, so it runs on L4s, up to two at a time while it scores shards.
# A warm pool of CPU workers takes the shard fan-out and the small global steps.
cpu_env = flyte.TaskEnvironment(
name="refinery-cpu",
image=cpu_image,
resources=flyte.Resources(cpu=2, memory="8Gi"),
reusable=flyte.ReusePolicy(replicas=(2, 8), concurrency=1, idle_ttl=600, scaledown_ttl=120),
include=CODE,
)
# The labeler: one L40S that keeps the vLLM server and the model weights loaded between calls.
label_env = flyte.TaskEnvironment(
name="refinery-labeler",
image=gpu_image,
resources=flyte.Resources(cpu=6, memory="40Gi", gpu="L40s:1", disk="60Gi"),
reusable=flyte.ReusePolicy(replicas=1, concurrency=1, idle_ttl=600, scaledown_ttl=300),
# vLLM's FlashInfer sampler compiles CUDA code at startup, which needs nvcc; the image has
# only the CUDA runtime, so use the PyTorch sampler.
env_vars={"VLLM_USE_FLASHINFER_SAMPLER": "0"},
include=CODE,
)
# The classifier is a 110M-parameter encoder, so a smaller GPU does. Two warm replicas split
# the scoring fan-out, and each keeps the encoder in memory between shards.
score_env = flyte.TaskEnvironment(
name="refinery-scorer",
image=gpu_image,
resources=flyte.Resources(cpu=4, memory="16Gi", gpu="L4:1", disk="40Gi"),
reusable=flyte.ReusePolicy(replicas=(1, 2), concurrency=1, idle_ttl=600, scaledown_ttl=180),
include=CODE,
)
driver_env = flyte.TaskEnvironment(
name="refinery-driver",
image=cpu_image,
resources=flyte.Resources(cpu=2, memory="4Gi"),
depends_on=[cpu_env, label_env, score_env],
include=CODE,
)
The stages hand each other small dataclasses. The documents themselves travel as Parquet
files (flyte.io.File), so a dataclass carries only a reference, counts, and a few examples
for the report:
@dataclass
class ShardSpec:
index: int
file: str
row_groups: list[int]
@dataclass
class EvalNgrams:
hashes: File # .npy of sorted uint64 13-gram hashes
items: File # parquet: benchmark, text, one row per eval question; ``hashes`` index into it
counts: dict[str, int] = field(default_factory=dict)
@dataclass
class ShardResult:
index: int
docs: File # parquet of the documents that passed the heuristics, with their signatures
raw: int
kept: int
rejected: dict[str, int] = field(default_factory=dict)
examples: list[dict] = field(default_factory=list)
contaminated: int = 0
contam_examples: list[dict] = field(default_factory=list)
seconds: float = 0.0
@dataclass
class DedupResult:
removed: File # parquet: id, stage, kept_id
after_exact: int
after_near: int
after_decontam: int
examples: list[dict] = field(default_factory=list)
clusters: list[dict] = field(default_factory=list)
seconds: float = 0.0
@dataclass
class LabelChunk:
labels: File # parquet: id, llm_score, justification, token counts
n: int
parsed: int
seconds: float # time spent labeling, without the model load
startup_s: float # vLLM start and model load, if this call paid for it
prompt_tokens: int
completion_tokens: int
@dataclass
class ClassifierResult:
head: File
held_out: dict # distilled classifier against held-out LLM labels
public_vs_llm: dict # public FineWeb-Edu classifier against the same held-out LLM labels
ours_vs_public: dict
points: list[list[float]] = field(default_factory=list) # held out: [llm, ours, public]
label_counts: list[int] = field(default_factory=list)
examples: list[dict] = field(default_factory=list)
n_train: int = 0
curve: list[float] = field(default_factory=list)
@dataclass
class ScoredShard:
scores: File # parquet: id, edu_score
n: int
seconds: float # time spent scoring, without the model load
load_s: float
@dataclass
class Curated:
kept: int
scored: int
tokens_kept: int
tokens_scored: int
histogram: list[int] = field(default_factory=list) # 20 bins of 0.25 over [0, 5]
examples: list[dict] = field(default_factory=list)
Heuristic filters, per shard
The input is HuggingFaceFW/fineweb, sample-10BT, read straight from the Hugging Face Hub as
Parquet row groups. Each shard is 20 row groups, about 20,000 documents, and the shards are
spread over the sample’s 14 files. The sample is drawn from every snapshot, so each shard
mixes documents from several Common Crawl snapshots (8 in each 20,000-document slice we
checked).
The filters are the document-level rules from Gopher, C4, and FineWeb, with their published
thresholds. They are a dataclass, and the dataclass is an input of every shard task, so
changing a threshold gives every shard a new cache key. The cache does not see edits to the
code in filters.py, because cache="auto" hashes only the task’s own source: when you change
a rule, bump FILTERS_VERSION, which is a field of the dataclass.
@dataclass
class FilterConfig:
"""Every threshold the CPU stages use. It is an input of each shard task, so changing any
value here gives new cache keys, and leaving it alone lets a rerun reuse every shard."""
# Gopher quality (Rae et al. 2021)
min_words: int = 50
max_words: int = 100_000
min_avg_word_len: float = 3.0
max_avg_word_len: float = 10.0
max_symbol_word_ratio: float = 0.1
max_bullet_lines_ratio: float = 0.9
max_ellipsis_lines_ratio: float = 0.3
max_non_alpha_words_ratio: float = 0.8
min_stop_words: int = 2
# Gopher repetition
max_dup_line_frac: float = 0.3
max_dup_para_frac: float = 0.3
max_dup_line_char_frac: float = 0.2
max_dup_para_char_frac: float = 0.2
# FineWeb's extra rules
min_line_punct_ratio: float = 0.12
max_short_line_ratio: float = 0.67
max_newline_word_ratio: float = 0.3
# MinHash: 14 bands of 8 hashes over word 5-grams, as in FineWeb
shingle: int = 5
bands: int = 14
rows: int = 8
seed: int = 1
# Decontamination
contam_ngram: int = 13
version: str = FILTERS_VERSION
def heuristic_reject(text: str, c: FilterConfig, words: list[str], hashes) -> str | None:
"""Return the name of the first rule the document fails, or None if it passes them all.
``words`` and ``hashes`` are the document's tokens and their hashes."""
raw_words = text.split()
n = len(raw_words)
if n < c.min_words:
return "gopher: too few words"
if n > c.max_words:
return "gopher: too many words"
avg = sum(len(w) for w in raw_words) / n
if not c.min_avg_word_len <= avg <= c.max_avg_word_len:
return "gopher: mean word length"
if (text.count("#") + text.count("...") + text.count("…")) / n > c.max_symbol_word_ratio:
return "gopher: symbol-to-word ratio"
lines = [ln.strip() for ln in text.split("\n")]
nonempty = [ln for ln in lines if ln]
if nonempty:
if sum(ln.startswith(BULLETS) for ln in nonempty) / len(nonempty) > c.max_bullet_lines_ratio:
return "gopher: bullet lines"
if sum(ln.endswith(("...", "…")) for ln in nonempty) / len(nonempty) > c.max_ellipsis_lines_ratio:
return "gopher: ellipsis lines"
if sum(not any(ch.isalpha() for ch in w) for w in raw_words) / n > 1 - c.max_non_alpha_words_ratio:
return "gopher: non-alphabetic words"
if len(STOP_WORDS.intersection(words)) < c.min_stop_words:
return "gopher: missing stop words"
# C4: placeholder text and code
low = text.lower()
if "lorem ipsum" in low:
return "c4: lorem ipsum"
if "{" in text:
return "c4: curly bracket"
# Gopher repetition: whole lines and paragraphs, then n-grams
paras = [p.strip() for p in re.split(r"\n{2,}", text) if p.strip()]
line_frac, line_chars = _dup_frac(nonempty)
para_frac, para_chars = _dup_frac(paras)
if line_frac > c.max_dup_line_frac or para_frac > c.max_dup_para_frac:
return "gopher: duplicate lines"
if line_chars > c.max_dup_line_char_frac or para_chars > c.max_dup_para_char_frac:
return "gopher: duplicate line characters"
reason = _ngram_repetition(words, hashes)
if reason:
return reason
# FineWeb's three rules
if nonempty:
if sum(ln.endswith(END_PUNCT) for ln in nonempty) / len(nonempty) < c.min_line_punct_ratio:
return "fineweb: lines without end punctuation"
if sum(len(ln) <= 30 for ln in nonempty) / len(nonempty) > c.max_short_line_ratio:
return "fineweb: short lines"
if text.count("\n") / n > c.max_newline_word_ratio:
return "fineweb: newline-to-word ratio"
return None
A document that passes gets the signatures the global steps need: a hash of its normalized text for exact duplicates, and a MinHash signature over word 5-grams cut into 14 bands of 8 for near duplicates, as in FineWeb.
def minhash(h, c: FilterConfig):
"""MinHash signature over word shingles: for each of bands*rows hash functions
(multiply-shift: (a*x + b) >> 32), the minimum over the document's shingles."""
import numpy as np
key = (c.seed, c.bands, c.rows)
if key not in _PERM_CACHE:
_PERM_CACHE[key] = _perms(c)
a, b = _PERM_CACHE[key]
shingles = np.unique(ngram_hashes(h, c.shingle))
with np.errstate(over="ignore"):
return ((a[:, None] * shingles[None, :] + b[:, None]) >> np.uint64(32)).min(axis=1)
def band_keys(sig, c: FilterConfig):
"""One 64-bit key per band of ``rows`` signature values. Two documents that share any band
key are near-duplicate candidates. With 14 bands of 8, a pair with Jaccard similarity s
becomes a candidate with probability 1 - (1 - s**8)**14: about 50% at s = 0.69, 77% at
0.75, 99% at 0.85, and about 5% at 0.5."""
import numpy as np
rows = sig.reshape(c.bands, c.rows)
key = np.arange(c.bands, dtype=np.uint64) # different bands never collide
with np.errstate(over="ignore"):
for j in range(c.rows):
key = key * np.uint64(_P) + rows[:, j]
return key
The shard task also checks every document against the eval sets. build_eval_ngrams hashes
every 13-gram of every GSM8K, MMLU, and ARC test question followed by its answer options, the
same way the shard task hashes documents. (The code also hashes the question alone, which adds
nothing: its 13-grams are a prefix of the longer text.)
EVAL_SETS = {
"gsm8k": ("openai/gsm8k", "main/test-00000-of-00001.parquet"),
"mmlu": ("cais/mmlu", "all/test-00000-of-00001.parquet"),
"arc-challenge": ("allenai/ai2_arc", "ARC-Challenge/test-00000-of-00001.parquet"),
"arc-easy": ("allenai/ai2_arc", "ARC-Easy/test-00000-of-00001.parquet"),
}
@cpu_env.task(cache="auto", retries=2)
async def build_eval_ngrams(config: filters.FilterConfig) -> EvalNgrams:
"""Every 13-gram of every test question in GSM8K, MMLU and ARC, hashed like document text."""
import numpy as np
import pandas as pd
rows = []
for name, (repo, path) in EVAL_SETS.items():
df = pd.read_parquet(f"https://huggingface.co/datasets/{repo}/resolve/main/{path}")
for q in df.to_dict("records"):
question = q["question"]
choices = q.get("choices")
if isinstance(choices, dict): # ARC: {"text": [...], "label": [...]}
choices = choices["text"]
extra = " ".join(map(str, choices)) if choices is not None else ""
rows.append((name, question, extra))
hashes, owners = [], []
for i, (_, question, extra) in enumerate(rows):
# The question alone, and the question followed by its answer options.
grams = np.unique(np.concatenate([
filters.ngram_hashes(filters.token_hashes(filters.tokens(t)), config.contam_ngram)
for t in (question, f"{question} {extra}")
]))
hashes.append(grams)
owners.append(np.full(len(grams), i, dtype=np.int64))
h, o = np.concatenate(hashes), np.concatenate(owners)
order = np.argsort(h, kind="stable")
d = tempfile.mkdtemp()
np.save(f"{d}/ngrams.npy", np.stack([h[order], o[order].astype(np.uint64)]))
pd.DataFrame(rows, columns=["benchmark", "text", "choices"]).to_parquet(f"{d}/items.parquet")
counts = pd.Series([r[0] for r in rows]).value_counts().to_dict()
print(f"{len(rows)} eval questions, {len(h)} distinct 13-grams")
return EvalNgrams(await File.from_local(f"{d}/ngrams.npy"), await File.from_local(f"{d}/items.parquet"),
{k: int(v) for k, v in counts.items()})
@cpu_env.task(cache="auto", retries=2)
async def refine_shard(spec: ShardSpec, config: filters.FilterConfig, evals: EvalNgrams) -> ShardResult:
"""Read one shard of raw web text, drop documents that fail a heuristic, and compute what the
global steps need: an exact-match hash, MinHash band keys, and any overlap with an eval set."""
import re
import fsspec
import numpy as np
import pandas as pd
import pyarrow.parquet as pq
started = time.monotonic()
url = f"https://huggingface.co/datasets/{DATASET}/resolve/main/{spec.file}"
with fsspec.open(url, "rb", block_size=16 * 2**20) as f:
table = pq.ParquetFile(f).read_row_groups(spec.row_groups, columns=["id", "url", "dump", "token_count", "text"])
eval_grams = np.load(await evals.hashes.download())
eval_items = pd.read_parquet(await evals.items.download())
rows, rejected, examples, contam_examples = [], {}, [], []
for doc in table.to_pylist():
text = doc["text"]
words = filters.tokens(text)
hashes = filters.token_hashes(words)
reason = filters.heuristic_reject(text, config, words, hashes)
if reason:
rejected[reason] = rejected.get(reason, 0) + 1
if rejected[reason] <= 2:
examples.append({"stage": "heuristics", "rule": reason, "url": doc["url"], "text": _snippet(text)})
continue
keys = filters.band_keys(filters.minhash(hashes, config), config)
# Decontamination: does any 13-gram of the document appear in an eval question?
grams = filters.ngram_hashes(hashes, config.contam_ngram)
pos = np.searchsorted(eval_grams[0], grams).clip(max=eval_grams.shape[1] - 1)
hits = np.flatnonzero(eval_grams[0][pos] == grams)
contam = ""
if len(hits):
item = eval_items.iloc[int(eval_grams[1][pos[hits[0]]])]
contam = item["benchmark"]
spans = [m.span() for m in re.finditer(r"\w+", text)]
first = min(hits[0], len(spans) - config.contam_ngram)
start, end = spans[first][0], spans[first + config.contam_ngram - 1][1]
contam_examples.append({
"url": doc["url"], "benchmark": contam, "question": _snippet(item["text"], 400),
"match": text[start:end], "matched_ngrams": int(len(np.unique(grams[hits]))),
"context": _snippet(text[max(0, start - 250):end + 250], 700),
})
rows.append({
"id": doc["id"], "url": doc["url"], "dump": doc["dump"], "token_count": doc["token_count"], "text": text,
"exact": filters.exact_hash(text), "contam": contam,
**{f"b{i}": k for i, k in enumerate(keys)},
})
out = os.path.join(tempfile.mkdtemp(), f"shard-{spec.index:03d}.parquet")
pd.DataFrame(rows).to_parquet(out, index=False)
seconds = time.monotonic() - started
print(f"shard {spec.index}: {table.num_rows} read, {len(rows)} kept, {len(contam_examples)} contaminated, {seconds:.0f}s")
return ShardResult(spec.index, await File.from_local(out), table.num_rows, len(rows), rejected, examples,
len(contam_examples), contam_examples[:20], round(seconds, 1))
The task is cached on its inputs: the shard’s location, the filter settings, and the eval n-grams. A rerun with the same settings skips it.
Global deduplication and decontamination
Deduplication has to see every shard at once. global_dedup reads only the hash columns,
drops exact duplicates, and then treats two documents as near duplicates when they share any
band key. Candidate pairs become edges of a graph, each connected component is one cluster,
and each cluster keeps its first document. Documents that survive both and overlap an eval set
are removed last.
@cpu_env.task(cache="auto")
async def global_dedup(shards: list[ShardResult], config: filters.FilterConfig) -> DedupResult:
"""Remove exact duplicates, then near duplicates, across every shard; then contaminated documents.
Two documents are near-duplicate candidates when they share any MinHash band key. Candidates
are joined into clusters (connected components), and each cluster keeps its first document."""
import numpy as np
import pandas as pd
import pyarrow.parquet as pq
from scipy.sparse import coo_matrix
from scipy.sparse.csgraph import connected_components
started = time.monotonic()
bands = [f"b{i}" for i in range(config.bands)]
frames = []
for s in shards:
t = pq.read_table(await s.docs.download(), columns=["id", "url", "exact", "contam", *bands]).to_pandas()
t["shard"] = s.index
frames.append(t)
df = pd.concat(frames, ignore_index=True)
n = len(df)
exact_dup = df.duplicated("exact", keep="first").to_numpy()
first_of_hash = df.groupby("exact")["id"].transform("first").to_numpy()
alive = np.flatnonzero(~exact_dup)
src, dst = [], []
for b in bands:
keys = df[b].to_numpy()[alive]
order = np.argsort(keys, kind="stable")
same = keys[order][1:] == keys[order][:-1]
src.append(alive[order[:-1][same]])
dst.append(alive[order[1:][same]])
src, dst = np.concatenate(src), np.concatenate(dst)
graph = coo_matrix((np.ones(len(src)), (src, dst)), shape=(n, n))
_, cluster = connected_components(graph, directed=False)
rep = pd.Series(np.arange(n)[alive], index=alive).groupby(cluster[alive]).transform("min")
near_dup = np.zeros(n, dtype=bool)
near_dup[alive] = rep.to_numpy() != alive
rep_of = np.arange(n)
rep_of[alive] = rep.to_numpy()
after_near = ~exact_dup & ~near_dup
contaminated = after_near & (df["contam"].to_numpy() != "")
stage = np.where(exact_dup, "exact duplicate", np.where(near_dup, "near duplicate",
np.where(contaminated, "contaminated", "")))
removed = pd.DataFrame({"id": df["id"], "stage": stage,
"kept_id": np.where(exact_dup, first_of_hash, df["id"].to_numpy()[rep_of])})
removed = removed[removed["stage"] != ""]
out = os.path.join(tempfile.mkdtemp(), "removed.parquet")
removed.to_parquet(out, index=False)
# Examples for the report: a few duplicate pairs, with the text of both sides.
pairs = pd.concat([removed[removed.stage == "exact duplicate"].head(3), removed[removed.stage == "near duplicate"].head(6)])
wanted = set(pairs["id"]) | set(pairs["kept_id"])
text = {}
for s in shards:
t = pq.read_table(await s.docs.download(), columns=["id", "url", "text"]).to_pandas()
for r in t[t["id"].isin(wanted)].itertuples():
text[r.id] = (r.url, r.text)
examples = [{"stage": p.stage, "url": text[p.id][0], "text": _snippet(text[p.id][1], 400),
"kept_url": text[p.kept_id][0], "kept_text": _snippet(text[p.kept_id][1], 400)}
for p in pairs.itertuples() if p.id in text and p.kept_id in text]
sizes = pd.Series(cluster[alive]).value_counts()
clusters = [{"size": int(c), "url": df["url"].iloc[int(alive[np.flatnonzero(cluster[alive] == k)[0]])]}
for k, c in sizes[sizes > 1].head(5).items()]
result = DedupResult(await File.from_local(out), int(n - exact_dup.sum()), int(after_near.sum()),
int((after_near & ~contaminated).sum()), examples, clusters, round(time.monotonic() - started, 1))
print(f"{n} in, {result.after_exact} after exact, {result.after_near} after near, {result.after_decontam} after decontamination")
return result
Labeling with an LLM
The labeler uses the prompt FineWeb-Edu used with Llama-3-70B-Instruct: an additive 0-5
rubric for educational value at primary and grade school level, a short justification, and
a final Educational score: N.
# The FineWeb-Edu annotation prompt (Penedo et al. 2024, and the
# HuggingFaceFW/fineweb-edu-classifier model card). Llama-3-70B-Instruct answered it for
# 450,000 FineWeb pages; here Qwen2.5-14B-Instruct answers it for a few thousand.
PROMPT = """Below is an extract from a web page. Evaluate whether the page has a high educational value and could be useful in an educational setting for teaching from primary school to grade school levels using the additive 5-point scoring system described below. Points are accumulated based on the satisfaction of each criterion:
- Add 1 point if the extract provides some basic information relevant to educational topics, even if it includes some irrelevant or non-academic content like advertisements and promotional material.
- Add another point if the extract addresses certain elements pertinent to education but does not align closely with educational standards. It might mix educational content with non-educational material, offering a superficial overview of potentially useful topics, or presenting information in a disorganized manner and incoherent writing style.
- Award a third point if the extract is appropriate for educational use and introduces key concepts relevant to school curricula. It is coherent though it may not be comprehensive or could include some extraneous information. It may resemble an introductory section of a textbook or a basic tutorial that is suitable for learning but has notable limitations like treating concepts that are too complex for grade school students.
- Grant a fourth point if the extract highly relevant and beneficial for educational purposes for a level not higher than grade school, exhibiting a clear and consistent writing style. It could be similar to a chapter from a textbook or a tutorial, offering substantial educational content, including exercises and solutions, with minimal irrelevant information, and the concepts aren't too advanced for grade school students. The content is coherent, focused, and valuable for structured learning.
- Bestow a fifth point if the extract is outstanding in its educational value, perfectly suited for teaching either at primary school or grade school. It follows detailed reasoning, the writing style is easy to follow and offers profound and thorough insights into the subject matter, devoid of any non-educational or complex content.
The extract:
{extract}.
After examining the extract:
- Briefly justify your total score, up to 100 words.
- Conclude with the score using the format: "Educational score: <total points>\""""
SCORE = re.compile(r"Educational score:\s*\**\s*(\d)")
def parse_score(answer: str) -> int | None:
"""The last "Educational score: N" in the answer, or None if the model did not give one."""
found = SCORE.findall(answer)
return min(int(found[-1]), 5) if found else None
vLLM runs as an OpenAI-compatible server inside the reusable container. The first call on the replica starts it; a later call with the same model finds it running.
def get_server(model: str, timeout_s: int = 1200) -> tuple[Server, bool]:
"""Start vLLM once per container. Returns the server and whether this call started it."""
global _server
import httpx
if _server is not None and _server.process.poll() is None:
if _server.model == model:
return _server, False
_server.process.terminate() # a different labeler was asked for: swap models
_server.process.wait(timeout=120)
cmd = [sys.executable, "-m", "vllm.entrypoints.cli.main", "serve", model,
"--port", str(PORT), "--max-model-len", "8192", "--gpu-memory-utilization", "0.90",
"--max-num-seqs", "256"]
log_path = "/tmp/vllm.log"
started = time.monotonic()
process = subprocess.Popen(cmd, stdout=open(log_path, "w"), stderr=subprocess.STDOUT)
while time.monotonic() - started < timeout_s:
if process.poll() is not None:
raise RuntimeError(f"vLLM exited with code {process.returncode}:\n{tail(log_path)}")
try:
if httpx.get(f"http://127.0.0.1:{PORT}/health", timeout=5).status_code == 200:
_server = Server(model, process, time.monotonic() - started, log_path)
print(f"vLLM ready in {_server.startup_s:.0f}s", flush=True)
return _server, True
except httpx.HTTPError:
pass
time.sleep(3)
process.kill()
raise TimeoutError(f"vLLM did not start in {timeout_s}s:\n{tail(log_path)}")
async def label_all(texts: list[str], model: str, concurrency: int = 256) -> list[dict]:
"""Send every document to the server at once (bounded), greedy decoding, and parse the score."""
import httpx
sem = asyncio.Semaphore(concurrency)
async with httpx.AsyncClient(base_url=f"http://127.0.0.1:{PORT}", timeout=600) as client:
async def one(text: str) -> dict:
body = {
"model": model,
"messages": [{"role": "user", "content": PROMPT.format(extract=text[:MAX_DOC_CHARS])}],
"temperature": 0.0,
"max_tokens": 300,
}
async with sem:
r = await client.post("/v1/chat/completions", json=body)
r.raise_for_status()
out = r.json()
answer = out["choices"][0]["message"]["content"]
return {
"llm_score": parse_score(answer),
"justification": answer.strip(),
"prompt_tokens": out["usage"]["prompt_tokens"],
"completion_tokens": out["usage"]["completion_tokens"],
}
return await asyncio.gather(*[one(t) for t in texts])
The driver splits the 4,000-document sample into four chunks and labels them one after
another. Each chunk is a cached task, so a rerun never pays for labels it already has, and a
failure costs at most the chunk in flight: a retry relabels that one chunk, not the sample. cache="auto" hashes the task’s own source, not the modules it calls, so the task
salts its cache with labeler.VERSION: change the prompt, bump the version, and the labels are
recomputed.
@label_env.task(cache=flyte.Cache("auto", salt=labeler.VERSION), retries=1)
async def label_chunk(chunk: File, model: str) -> LabelChunk:
"""Score a chunk of documents on the FineWeb-Edu rubric with the LLM. The first call on a
replica starts vLLM; the rest find it warm."""
import pandas as pd
server, started_here = labeler.get_server(model)
df = pd.read_parquet(await chunk.download())
t0 = time.monotonic()
answers = await labeler.label_all(df["text"].tolist(), model)
seconds = time.monotonic() - t0
out = pd.DataFrame(answers)
out.insert(0, "id", df["id"].to_numpy())
path = os.path.join(tempfile.mkdtemp(), "labels.parquet")
out.to_parquet(path, index=False)
print(f"{len(df)} documents in {seconds:.0f}s, {out.llm_score.notna().sum()} with a score")
return LabelChunk(await File.from_local(path), len(df), int(out.llm_score.notna().sum()), round(seconds, 2),
round(server.startup_s, 1) if started_here else 0.0,
int(out.prompt_tokens.sum()), int(out.completion_tokens.sum()))
Distilling the labels into a classifier
The classifier follows the FineWeb-Edu classifier’s recipe: the
snowflake-arctic-embed-m encoder
stays frozen, and only a small head on its [CLS] vector learns to predict the LLM’s score.
Because the encoder is frozen, each document is encoded once and the head trains in seconds.
Web text is mostly scores 0 to 2, so the loss weights each document by the inverse frequency of
its score; without that, the head learns to predict “about 1” for everything.
def _batches(tok, texts: list[str], batch_size: int):
"""Tokenize, then batch documents of similar length together so padding stays small."""
enc = tok(texts, truncation=True, max_length=MAX_TOKENS)
order = sorted(range(len(texts)), key=lambda i: len(enc["input_ids"][i]))
for s in range(0, len(order), batch_size):
idx = order[s:s + batch_size]
yield idx, tok.pad({"input_ids": [enc["input_ids"][i] for i in idx],
"attention_mask": [enc["attention_mask"][i] for i in idx]}, return_tensors="pt")
def embed(texts: list[str], batch_size: int = 128):
"""The encoder's last-layer [CLS] vector for every text, in input order (float32, N x 768)."""
import numpy as np
import torch
tok, model, _ = _load(ENCODER)
out = np.zeros((len(texts), model.config.hidden_size), dtype=np.float32)
with torch.inference_mode():
for idx, batch in _batches(tok, texts, batch_size):
hidden = model(**{k: v.cuda() for k, v in batch.items()}).last_hidden_state[:, 0]
out[idx] = hidden.float().cpu().numpy()
return out
def make_head(dim: int = 768):
from torch import nn
return nn.Sequential(nn.Linear(dim, dim), nn.Tanh(), nn.Linear(dim, 1))
def train_head(x, y, epochs: int = 100, lr: float = 3e-4, batch_size: int = 256, seed: int = 0):
"""Fit the head to the LLM scores with a class-balanced MSE.
Web text is mostly scores 0-2, so a plain MSE learns to predict "about 1" for everything and
never reaches the threshold. Weighting each document by the inverse frequency of its score
(capped at 20x) makes the rare high scores count. 10% of the rows pick the best epoch."""
import numpy as np
import torch
torch.manual_seed(seed)
rng = np.random.default_rng(seed)
counts = np.bincount(np.rint(y).astype(int), minlength=6).astype(float)
weight = np.minimum(counts.max() / np.maximum(counts, 1), 20.0)[np.rint(y).astype(int)]
perm = rng.permutation(len(x))
n_val = max(len(x) // 10, 1)
val, tr = perm[:n_val], perm[n_val:]
def t(a, idx):
return torch.tensor(a[idx], dtype=torch.float32).cuda()
xt, yt, wt = t(x, tr), t(y, tr), t(weight, tr)
xv, yv, wv = t(x, val), t(y, val), t(weight, val)
head = make_head(x.shape[1]).cuda()
opt = torch.optim.AdamW(head.parameters(), lr=lr, weight_decay=0.05)
best, best_state, curve = float("inf"), None, []
for epoch in range(epochs):
head.train()
for s in torch.randperm(len(xt)).split(batch_size):
loss = (wt[s] * (head(xt[s]).squeeze(-1) - yt[s]) ** 2).sum() / wt[s].sum()
opt.zero_grad()
loss.backward()
opt.step()
head.eval()
with torch.no_grad():
v = ((wv * (head(xv).squeeze(-1) - yv) ** 2).sum() / wv.sum()).item()
curve.append(v)
if v < best:
best, best_state = v, {k: t.detach().clone() for k, t in head.state_dict().items()}
head.load_state_dict(best_state)
return head.cpu(), curve
train_classifier holds out 20% of the labeled documents by a hash of their ID, trains the
head on the rest, and scores the held-out documents with both the new classifier and the
public FineWeb-Edu classifier, so the report can compare all three.
@score_env.task(cache=flyte.Cache("auto", salt=classifier.VERSION), retries=1)
async def train_classifier(chunks: list[File], labels: list[LabelChunk], held_out_share: float = 0.2) -> ClassifierResult:
"""Distill the LLM's labels into the encoder + head classifier, and measure it on held-out
labels next to the public FineWeb-Edu classifier."""
import pandas as pd
import torch
docs = pd.concat([pd.read_parquet(await c.download()) for c in chunks])
lab = pd.concat([pd.read_parquet(await l.labels.download()) for l in labels])
df = docs.merge(lab, on="id").dropna(subset=["llm_score"]).reset_index(drop=True)
df["llm_score"] = df["llm_score"].astype(float)
# A stable split by document id, so a rerun holds out the same documents.
test = df["id"].map(lambda s: zlib.crc32(s.encode()) % 1000 < held_out_share * 1000).to_numpy()
x = classifier.embed(df["text"].tolist())
head, curve = classifier.train_head(x[~test], df["llm_score"].to_numpy()[~test])
held = df[test].reset_index(drop=True)
held["ours"] = classifier.predict(head, x[test])
held["public"] = classifier.public_scores(held["text"].tolist())
llm, ours, public = held["llm_score"].to_numpy(), held["ours"].to_numpy(), held["public"].to_numpy()
path = os.path.join(tempfile.mkdtemp(), "head.pt")
torch.save(head.state_dict(), path)
examples = [{"llm": int(r.llm_score), "ours": float(r.ours), "public": float(r.public),
"text": _snippet(r.text, 500), "justification": r.justification}
for score in range(6) for r in held[held["llm_score"] == score].head(1).itertuples()]
return ClassifierResult(
head=await File.from_local(path),
held_out=classifier.agreement(llm, ours),
public_vs_llm=classifier.agreement(llm, public),
ours_vs_public=classifier.agreement(public, ours),
points=[[float(a), round(float(b), 3), round(float(c), 3)] for a, b, c in zip(llm, ours, public)],
label_counts=[int((df["llm_score"] == s).sum()) for s in range(6)],
examples=examples,
n_train=int((~test).sum()),
curve=[round(v, 4) for v in curve],
)
def agreement(reference, predicted, threshold: float = 3.0) -> dict:
"""How well ``predicted`` scores track ``reference`` scores (both 0-5, either may be float)."""
import numpy as np
ref = np.asarray(reference, dtype=np.float64)
pred = np.asarray(predicted, dtype=np.float64)
def ranks(v):
order = np.argsort(v, kind="stable")
r = np.empty(len(v))
r[order] = np.arange(len(v))
for val in np.unique(v): # average ranks over ties (scores are mostly integers)
tie = v == val
r[tie] = r[tie].mean()
return r
ri, pi = np.clip(np.rint(ref), 0, 5).astype(int), np.clip(np.rint(pred), 0, 5).astype(int)
tp = int(((ref >= threshold) & (pred >= threshold)).sum())
fp = int(((ref < threshold) & (pred >= threshold)).sum())
fn = int(((ref >= threshold) & (pred < threshold)).sum())
precision = tp / (tp + fp) if tp + fp else 0.0
recall = tp / (tp + fn) if tp + fn else 0.0
confusion = np.zeros((6, 6), dtype=int)
np.add.at(confusion, (ri, pi), 1)
return {
"n": int(len(ref)),
"spearman": float(np.corrcoef(ranks(ref), ranks(pred))[0, 1]),
"pearson": float(np.corrcoef(ref, pred)[0, 1]),
"mae": float(np.abs(ref - pred).mean()),
"exact": float((ri == pi).mean()),
"within_one": float((np.abs(ri - pi) <= 1).mean()),
"precision": precision,
"recall": recall,
"f1": 2 * precision * recall / (precision + recall) if precision + recall else 0.0,
"confusion": confusion.tolist(), # rows: reference score, columns: predicted score
}
Then every surviving document is scored, one cached task per shard on the L4 pool:
@score_env.task(cache=flyte.Cache("auto", salt=classifier.VERSION), retries=1)
async def score_shard(docs: File, removed: File, head: File) -> ScoredShard:
"""Score every surviving document in one shard with the distilled classifier."""
import pandas as pd
import torch
df = await _survivors(docs, removed, ["id", "text"])
model = classifier.make_head()
model.load_state_dict(torch.load(await head.download(), weights_only=True))
load_s = classifier.warm_up() # zero when this replica already holds the encoder
t0 = time.monotonic()
scores = classifier.predict(model, classifier.embed(df["text"].tolist()))
seconds = time.monotonic() - t0
path = os.path.join(tempfile.mkdtemp(), "scores.parquet")
pd.DataFrame({"id": df["id"].to_numpy(), "edu_score": scores.clip(0, 5)}).to_parquet(path, index=False)
print(f"{len(df)} documents in {seconds:.0f}s")
return ScoredShard(await File.from_local(path), len(df), round(seconds, 2), round(load_s, 1))
Curating and publishing
curate joins the scores to the surviving documents and keeps those at or above the
threshold. It is the only stage whose inputs include the threshold.
@cpu_env.task(cache="auto")
async def curate(shards: list[ShardResult], dedup: DedupResult, scores: list[ScoredShard], threshold: float) -> tuple[DataFrame, Curated]:
"""Join the scores to the surviving documents and keep those at or above the threshold."""
import numpy as np
import pandas as pd
kept, hist, scored, tokens_scored = [], np.zeros(20, dtype=int), 0, 0
low = []
for s, sc in zip(shards, scores):
df = await _survivors(s.docs, dedup.removed, ["id", "url", "dump", "token_count", "text"])
df = df.merge(pd.read_parquet(await sc.scores.download()), on="id")
hist += np.histogram(df["edu_score"], bins=20, range=(0, 5))[0]
scored += len(df)
tokens_scored += int(df["token_count"].sum())
below = df[df["edu_score"] < threshold]
low.append(below.sample(n=min(2, len(below)), random_state=s.index))
kept.append(df[df["edu_score"] >= threshold])
out = pd.concat(kept, ignore_index=True)
out["edu_int_score"] = out["edu_score"].round().clip(0, 5).astype(int)
examples = [{"stage": "classifier", "rule": f"score {r.edu_score:.2f}", "url": r.url, "text": _snippet(r.text)}
for r in pd.concat(low).head(4).itertuples()]
examples += [{"stage": "kept", "rule": f"score {r.edu_score:.2f}", "url": r.url, "text": _snippet(r.text)}
for r in out.nlargest(3, "edu_score").itertuples()]
stats = Curated(len(out), scored, int(out["token_count"].sum()), tokens_scored, hist.tolist(), examples)
return await DataFrame.from_local(out), stats
The driver runs the stages, updates the report as each one finishes, and publishes the result
as the edu-web-corpus artifact. The card records the funnel, every filter threshold, the
labeler, and the classifier’s agreement, so anyone who loads the dataset can see how it was
made.
@driver_env.task(report=True, produces_artifacts=True)
async def refine(
n_shards: int = 15,
docs_per_shard: int = 20_000,
n_labels: int = 4000,
label_chunks: int = 4,
threshold: float = 3.0,
label_model: str = "Qwen/Qwen2.5-14B-Instruct",
) -> DataFrame:
"""Turn raw web text into a scored, filtered, deduplicated and decontaminated corpus."""
config = filters.FilterConfig()
timing: dict[str, float] = {}
t = time.monotonic()
with flyte.group("filter-shards"):
evals = await build_eval_ngrams(config)
shards = list(await asyncio.gather(*[refine_shard(s, config, evals) for s in shard_specs(n_shards, docs_per_shard)]))
timing["filter shards"] = time.monotonic() - t
await flyte.report.replace.aio(reports.live(config, shards, None, None, None, None, threshold, evals.counts), do_flush=True)
t = time.monotonic()
dedup = await global_dedup(shards, config)
timing["dedup + decontaminate"] = time.monotonic() - t
await flyte.report.replace.aio(reports.live(config, shards, dedup, None, None, None, threshold, evals.counts), do_flush=True)
t = time.monotonic()
with flyte.group("llm-labels"):
chunks = await sample_for_labeling(shards, dedup, n_labels, label_chunks)
labels = []
for c in chunks: # one warm replica: the first call loads the model, the rest reuse it
labels.append(await label_chunk(c, label_model))
timing["LLM labels"] = time.monotonic() - t
t = time.monotonic()
clf = await train_classifier(chunks, labels)
timing["train classifier"] = time.monotonic() - t
await flyte.report.replace.aio(reports.live(config, shards, dedup, labels, clf, None, threshold, evals.counts), do_flush=True)
t = time.monotonic()
with flyte.group("score-shards"):
scores = list(await asyncio.gather(*[score_shard(s.docs, dedup.removed, clf.head) for s in shards]))
timing["score every document"] = time.monotonic() - t
t = time.monotonic()
dataset, curated = await curate(shards, dedup, scores, threshold)
timing["curate"] = time.monotonic() - t
summary = reports.summarize(config, shards, dedup, labels, clf, scores, curated, threshold, timing, label_model)
await flyte.report.replace.aio(reports.final(summary, shards, dedup, clf, curated, evals.counts), do_flush=True)
flyte.report.get_tab("Rejected documents").replace(reports.rejected_tab(shards, dedup, curated))
flyte.report.get_tab("LLM labels and the classifier").replace(reports.labels_tab(clf, labels))
await flyte.report.flush.aio()
card = await artifacts.Card.create_from.aio(content=reports.card(summary), format="html", card_type="data")
metadata = artifacts.Metadata(
name="edu-web-corpus",
description=f"{curated.kept:,} FineWeb documents with an educational score of at least {threshold}, "
f"from {summary['funnel'][0][1]:,} raw documents",
kind="data",
attrs={"threshold": str(threshold), "filters": config.version, "labeler": label_model,
"classifier": f"{classifier.ENCODER} + head", "documents": str(curated.kept),
"tokens": str(curated.tokens_kept), "source": f"{DATASET} sample-10BT"},
card=card,
)
return artifacts.new(dataset, metadata)
Results
A run with the defaults read 299,906 FineWeb documents in 15 shards. The CPU stages took about 5 minutes for the shards and 20 seconds for the global step. Labeling 4,000 documents with Qwen2.5-14B-Instruct took 12 minutes on one L40S, plus 6 minutes to download and load the model on the first call. Training the head took 41 seconds and scoring every surviving document took 7 minutes on the L4 pool. A cold run takes about 30 minutes.
| Stage | Documents kept | Share of raw |
|---|---|---|
| Raw FineWeb documents | 299,906 | 100% |
| Pass the heuristic filters | 290,179 | 96.8% |
| No exact duplicate | 290,159 | 96.7% |
| No near duplicate | 290,100 | 96.7% |
| No eval overlap | 290,068 | 96.7% |
| Educational score of at least 3 | 3,454 | 1.2% (3.8M tokens) |
| Educational score of at least 2.5 | 17,967 | 6.0% (19.0M tokens) |
The early stages remove little because FineWeb already ran the same kind of pipeline: Trafilatura extraction, language identification, Gopher and C4 filters, and MinHash deduplication within each Common Crawl snapshot. The rules that still fire most often are Gopher’s repeated 5-grams (2,682 documents) and its minimum length (1,827). Many of the 79 duplicates are the same URL crawled in two snapshots, which FineWeb deduplicated within a snapshot but not across them; the rest are templated pages. On raw crawl, the heuristic and deduplication stages remove far more; on FineWeb, as in FineWeb-Edu, the quality classifier does most of the filtering.
The 32 documents removed as eval overlap are worth reading. The ones in the report match MMLU questions that quote a famous public text, such as the Declaration of Independence, a Patrick Henry speech, or the boilerplate of a US statute. A 13-gram match finds exactly these documents, and whether to drop them is a policy decision: they contain the source passage, not the question or its answer.
Why distill
| GPU time per document | GPU | |
|---|---|---|
| LLM labeler, Qwen2.5-14B-Instruct | 174 ms | L40S |
| Distilled classifier | 2.6 ms | L4 |
The classifier is 66 times cheaper per document, on a smaller GPU. Labeling all 290,068 surviving documents with the LLM would have taken about 14 L40S-hours; the classifier took 13 L4-minutes. At FineWeb’s scale of billions of documents, only the classifier is affordable.
What it costs in agreement, on 797 held-out documents that the LLM labeled. The decision that matters for curation is the binary one, keep or drop at a score of 3, so F1 at that threshold comes first:
| Comparison | F1 at a score of 3 | Spearman | Within 1 point |
|---|---|---|---|
| Distilled classifier against the LLM labels | 0.20 | 0.61 | 98% |
| Public FineWeb-Edu classifier against the LLM labels | 0.31 | 0.65 | 98% |
| Distilled classifier against the public classifier | 0.21 | 0.74 | 99% |
At the threshold, the distilled head is clearly worse than the public classifier, which was
trained on 450,000 labels from a 70B model; on ranking, the gap is smaller. Read these numbers
with two caveats. The held-out labels come from the distilled head’s own teacher, so they
measure how well it learned the 14B model, not how good the 14B model is. And both classifiers
put a head on the same frozen arctic-embed-m encoder, which inflates their agreement with
each other (0.74). Only 43 held-out documents have an LLM score of 3 or more, so the F1 values
are noisy too.
The distilled scores are also compressed toward the middle: the head rarely predicts above 3.5, so at FineWeb-Edu’s threshold of 3 it keeps 1.2% of the corpus, while the LLM gave 4.8% of its sample a 3 or more. A threshold of 2.5 keeps 6.0%.
Changing the threshold
Rerunning with --threshold 2.5 took 33 seconds. 38 of the 39 tasks were cache hits: every
shard, the global step, the labels, the classifier, and every score. Only curate ran, and
the artifact got a new version with 17,967 documents.
What the iterations taught us
The first full run used Qwen2.5-3B-Instruct as the labeler. It was fast, 47 ms per document, but lenient: it gave 54% of the sample a 3 or more, against 5% for the 14B model. The public classifier agreed with its labels at an F1 of 0.08 at the threshold, and the classifier distilled from them kept 17% of the corpus. A distilled classifier is only as good as its labels, and its agreement with its own teacher hides that. Check the teacher against an independent reference, here the public classifier, before you distill it.
With the 14B labels, the first head was trained with a plain MSE. It ranked well, at a Spearman of 0.64, but almost every prediction fell between 0 and 2.5, and the threshold of 3 kept 30 documents. Weighting the loss by inverse score frequency spread the predictions out at a small cost in rank correlation.
At real scale
- More data. The shard task is the unit of work. Point
shard_specsat more files and raise the pool’sreplicas; the global step reads only the hash columns, so it stays small. Past a few hundred million documents, run deduplication as a distributed join instead. - More labels, a stronger labeler. FineWeb-Edu used 450,000 labels from a 70B model. Raise
n_labelsandlabel_chunks, and givelabel_envmore replicas or GPUs; every chunk is cached, so a long labeling job survives a failure. - A raw crawl. Replace the FineWeb reader with text extracted from Common Crawl WARC files and add a language identification stage. The heuristic and deduplication stages then do much more of the filtering.
- Other labels. Swap the prompt for a taxonomy, as Essential-Web did, and train one head per category on the same frozen embeddings.