Repository navigation
Add Mitra-v2 regressor via optional mitra-finetune dependency - #295
Merged
Merged
Conversation
Add FoundationModelConfig.optional_dependencies, pip-installed in the training / serving container before the model is used. mitra-v2-regressor ships mitra-finetune==0.3.0 (TestPyPI), whose patches switch AutoGluon 1.6 Mitra to the v2 cross-entropy regression head.
… pip in the scripts Rename FoundationModelConfig.optional_dependencies to requirements. The backend writes them as requirements.txt into the training sourcedir and the model's code/ dir, which the DLC installs before running the script.
…Mitra patch helper instead of duplicating it
…rparameters, add Mitra-v2 regressor docs and tests
shchur
marked this pull request as ready for review
October 12, 2026 14:13
shchur
changed the title
shchur
enabled auto-merge (squash)
October 12, 2026 14:18
|
Job PR-295-daccfdd is done. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds
mitra-v2-regressor, which runs through the patches in mitra-finetune 0.3.0. The patch is temporary: AutoGluon 1.7 is expected to support the checkpoint natively.FoundationModelConfig.requirements: list[str]: lines of arequirements.txtthat the container installs before the model is used. The backend writes it into the trainingsourcedir.tar.gz(installed bysagemaker_training) and into the model'scode/dir (installed by the DLCserveentrypoint'sinstall_requirements.pybefore Gunicorn starts). Both honourCA_REPOSITORY_ARN. Coversfit()/ batch predict, freshdeploy()andcache_model_artifact().mitra-v2-regressorregistry entry (autogluon/mitra-regressor-2) withrequirements=["mitra-finetune==0.3.0"].serving_utils/mitra.py:patch_mitra_regressorappliesinstall_d2h_sync_patch,install_reg_ce_patches(n_bins)(n_bins from checkpointconfig.json) andinstall_use_hf_patchifmitra_finetuneis importable andhf_reg_modelis set. Only the v2 registry entry installs the package, so nothing else is patched.serving_utils/now also ships in the training sourcedir, sotabular_fm_serve.pyandtrain.pyshare the helper.mitra-v2-regressorrow in the tabular foundation model table.mitra-v2-regressor; training sourcedir contents with and withoutrequirements.txt.Caveat: custom images not built on the DLC entrypoint won't install
requirements.txt.