[Draft] Make processing service support model retraining

[mohamedelabbas1996]

Summary

Antenna can now retrain a classifier head from species that people have verified, and this is the half that does the training. A head sits on a frozen backbone, so retraining it is fitting a small matrix over embeddings that already exist rather than another pass over the images — seconds, not weeks.

Antenna prepares the training set, writes it to storage and hands over a URL. This downloads it, fits a new head, scores it against the head currently in service on the same held-out rows, saves it, and offers it as its own pipeline. It never replaces the running head: promoting one is a separate, deliberate step, because an automatic swap would let a single bad run quietly degrade every later classification.

The Antenna side is RolnickLab/antenna#1407.

List of Changes

#	What it does	How
1	A model can declare that its head can be retrained	trainable on InferenceBaseClass, defaulting to False, mirrored into the algorithm config Antenna reads
2	Serves BioCLIP 2.5 classifiers whose head can be swapped	trapdata/ml/models/bioclip.py: frozen ViT-H/14 backbone, linear head read from the Hub or from disk
3	A classification carries the embedding that produced it	features on the classification response, so a caller can keep the vector without running the backbone twice
4	Retrains a head from a prepared dataset	POST /train, and trapdata/ml/training.py
5	Keeps species the old head knew but the new data does not cover	warm_start and restore_untrained_classes
6	Scores the new head against the one in service	score_incumbent and evaluate, on the same held-out rows
7	Every head it produces becomes selectable	trapdata/ml/trained_heads.py registers each saved head as its own pipeline, without a restart
8	Reports the outcome back to the Antenna job	signed-token callback, so training can outlast the request that started it
9	Uploads the head so it outlives this service's disk	optional upload URL in the request; the response says where it landed

Detailed Description

Why the dataset arrives as a URL

A project with a few hundred thousand verified labels runs to hundreds of megabytes. Sending that in one request is fragile and has to start over if the connection drops. A URL is small to hand over and cheap to retry.

Why nothing is promoted automatically

/train reports promote and the reason for it, and stops there. The head it was trained from stays exactly where it was; a retrain adds a choice rather than replacing one. A held-out set under MIN_MEANINGFUL_TEST_ROWS (30) is reported as too small to tell two heads apart, so a smoke test is not mistaken for evidence.

Why the species list comes from the dataset

Antenna sends the project's taxa list inside the dataset, and it wins over whatever happens to be in the rows. A species with no verified crops yet must still be something the head can predict, or the pipeline goes blind to it the moment it is retrained.

Why the head is uploaded

It used to be written only to this service's own disk, under a cache directory, so a cleared cache or a rebuilt container lost it while Antenna still reported the version as existing. The upload is optional: a caller that sends no URL gets exactly the previous behaviour, and a failed upload is logged rather than raised, since the training succeeded and the head is still on local disk.

Verification

• test_head_retraining.py, test_bioclip_classifier.py and test_trainable_flag.py cover fitting, the declared species list winning over the rows, an unsupported head type being refused, a saved head being discoverable and servable, the three label-map shapes, and the upload including its failure path.
• Run end to end against a live Antenna: a job started from the UI prepared 96 verified crops over 17 species, this service trained and saved a head, uploaded it, and Antenna registered the new version with the location recorded.

Known gaps

1. A stored head is never fetched back. The weights are durable and their location is recorded, but a service that has lost its disk copy does not re-download them; that needs head discovery at startup.
2. The incumbent is looked up by key. A head registered under a different key than the one Antenna asks for is reported as having no baseline, so promote can only ever be false for it. Worth deciding whether to fall back to the parent recorded in the training info.
3. A classifier with neither a local head directory nor a published repository crashes the service at startup, with an error from the Hub client rather than a clear message. It should skip that pipeline and say why.
4. Only linear heads are supported. Anything else is refused rather than quietly served as a linear head.

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}