SetFit, but with images

A few weeks ago Daniel van Strien wrote up how he used agents to build a classifier for data curation: an agent labels a couple of hundred PDFs under a human-reviewed guide, SetFit trains a small classifier, Hugging Face Jobs runs it over 191,000 documents, everyone is home before dinner.

One thought would not leave me alone. SetFit is a text method. My documents are PDFs, often scanned, and what tells you a PDF is an earnings release rather than an annual report is usually the shape of the first page, not the words. Could SetFit do the same trick on images?

The rest is history. Also about $2 of GPU, one pull request, and a lesson about labels.

SetFit is a two-trick pony, and neither trick cares about text

SetFit does two things: fine-tune a sentence-transformer on pairs (same class pulls together, different class pushes apart), then fit a logistic regression on the embeddings. If you are a careful reader, you must have noticed nothing in that says "text". It needs a sentence-transformer that can embed your inputs, and for page images NeoMME is one (thank you Tony Wu and Aurélien Lac) Smaller version is a 260M parameters multimodal encoder. Therefore a PIL image in and a vector out, with a task="document" switch you must remember to flip.

SetFit did assume strings in a few corners (the pair sampler, the model card, routing that task argument), so I patched them: huggingface/setfit#653 (Not yet in main as of 9/11). With it:

from setfit import SetFitModel, Trainer, TrainingArguments

model = SetFitModel.from_pretrained("Hcompany/NeoMME-260M-Retriever-ST-dense", task="document")
Trainer(model=model, train_dataset=pages,  # PIL images in the "text" column; naming things is hard
        args=TrainingArguments(num_iterations=5, batch_size=4)).train()
model.predict([Image.open("page-1.jpg")])

It worked on my financial filings with good accuracy and F1. But my data and my labels prove little, so I went back to the source.

Daniel's documents, Daniel's labels

Daniel published his 1% sample with his classifier's predictions (davanstrien/finepdfs-edu-purpose), six purposes: administration_policy, exercise_assessment, instruction_reference, news_promotion, research_analysis, other. I fetched the PDFs from their original URLs, rendered page one with pdftoppm, and trained on his labels with an image twin of the uv-scripts SetFit recipe:

hf jobs uv run --flavor t4-small --secrets HF_TOKEN \
  train-setfit-images.py oneryalcin/finepdfs-purpose-pages my-org/finepdfs-neomme --num-samples 64

Thirty-three minutes on a T4 later: 0.55 accuracy against his labels on 703 held-out pages. Above chance (0.2), well below the text model. Exam sheets and research papers it got; "instruction" versus "administration" from a cover page, not so much.

I nearly wrote this up as "text wins, images are cute". Then I asked the boring question.

Better according to whom?

Those labels are a model's predictions, and Daniel says so plainly: 65.8% accuracy against agent references on his own held-out set. So my 0.55 was agreement with something that is right about two thirds of the time. When the reference is a model, a second model disagreeing with it is not automatically wrong. You need a referee.

I then what Daniel did, hired DeepSeek-V4.1-Flash. It takes images, costs $0.15 per million input tokens, and has no stake in the outcome. Each document went in as the page image plus a start/middle/end text excerpt plus Daniel's purpose guide, and came back as one label and a twenty-word reason:

GUIDE = """Label the PRIMARY PURPOSE of this PDF. One of:
exercise_assessment (worksheets, exams, answer keys) · instruction_reference (teaching material, guides,
reference) · research_analysis (papers, investigations, analytical essays) · administration_policy
(rules, requirements, schedules, business records) · news_promotion (newsletters, announcements,
promotional material) · other.
Answer with JSON only: {"label": "...", "reason": "<= 20 words"}"""

The judge is an LLM. It has opinions. But they are different opinions from both contestants, which is the job. Two things fell out. It agrees with Daniel's classifier on 68% of documents, which is his own number reproduced. And on the 313 documents where text and image disagreed, it sided with text 165 times, with the image 70, and with neither 78. The text model really was better on those labels; the image model had been faithfully learning a third noise.

If a third of your training labels are wrong, do not tune the model. Fix the labels.

Same models, clean labels

Relabel everything with the judge: 4,972 training pages, 702 test pages, about $3. Retrain. And this time also keep the frozen version, which is two lines and needs no GPU:

emb = body.encode(images, task="document", normalize_embeddings=True)   # 5 min on a laptop GPU
head = LogisticRegression(C=100, max_iter=5000).fit(emb_train, y_train)

Same 702 documents, same judge labels, every row:

model	input	accuracy	macro-F1
Daniel's SetFit (ModernBERT-embed)	text excerpt	0.679	0.570
NeoMME, frozen + logistic head	page-1 image	0.694	0.634
NeoMME, contrastive fine-tune	page-1 image	0.681	0.638

Honest reading: on par. Error bars on 702 documents are about ±3.5 points, the judge is an LLM, and I picked the head's C on this same split, which is the kind of thing you confess in a blog post so nobody finds it in code review. What I will claim: a classifier that never extracts a character of text lands where a text classifier does, on a corpus that is mostly about content, from five thousand judge-labelled pages and a laptop.

The fine-tune, 729 steps and $1.70 of H100, did not beat the frozen head. It had not on my French filings either (0.763 vs 0.759). The pony's first trick seems to be for eight examples, not five thousand.

Both model and dataset is published in hf:
dataset oneryalcin/finepdfs-purpose-pages-v2,

models finepdfs-purpose-neomme-frozen-head and finepdfs-purpose-setfit-neomme-v2.

Things I would tell past me

The whole gain, 0.55 to 0.69, came from relabelling and more data. The model never changed. Labels beat modality, and three dollars of judge beat any hyper-parameter I could have turned. Which leads to the second thing: everyone evaluates against a judge, far fewer people train against one, and that is where the money is.

Run the frozen version first. Embeddings plus a logistic regression takes five minutes on a laptop and tells you whether the model can see your problem at all. Twice now the contrastive fine-tune has failed to beat it.

Two engineering notes, since they cost me an evening. SetFit's pair sampler is quadratic and stores your inputs in the pair table. Harmless for strings. With 3,887 images (be honest I'm stretching setfits proposal of few lables here, so my mistake in the firts place) it built 7.5 million candidate pairs and tried to write 78,000 page images into one Arrow table, which died after an hour of H100 with ArrowInvalid: offset overflow. The PR now stores indices and lets the collator look images up at batch time; over an hour became 0.4 seconds. And feed the GPU some CPUs: image preprocessing lives in the dataloader, and Modal's default single core left an H100 at 0% utilisation, looking very expensive and very idle. Eight cores and eight workers pinned it at 99%.

Open ends

Daniel mentions 320 documents labelled under a human-reviewed guide. If those ever go public they are a far better test set than my LLM referee, and I would happily rerun all three rows. Multi-page input, image plus text, and non-English corpora are the obvious next experiments. And if you train SetFit on anything that is not a string, the PR could use your eyes.

Thanks to Daniel for a post clear enough to argue with. Highest compliment I know.