Deploy MLPMetricFull v2 (47k models, with ID emb)

.gitattributes

@@ -1,35 +1,2 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text

README.md

@@ -15,9 +15,11 @@ short_description: Finding the Best Model for Your Task from Myriads of Models
 # ModelLens — Finding the Best Model for Your Task from Myriads of Models
 
 Describe your dataset → pick a task and metric → get a ranked list of HuggingFace
-models likely to perform well on it. Backed by the `MLPMetric` (ablation_no_id)
-checkpoint trained on the `unified_augmented` corpus, with a candidate pool of
-~47k HuggingFace models.
+models likely to perform well on it. Backed by the `MLPMetricFull` checkpoint
+trained on the cleaned + expanded `unified_augmented_v2` corpus, with a candidate
+pool of ~47k HuggingFace models. The full model uses learned model-id /
+model-description / dataset-id embeddings on top of the dataset-description and
+task/metric signals.
 
 ## How it works
 
@@ -49,16 +51,17 @@ requirements.txt    Pinned deps
 assets/
   model_pool.npz    Pre-computed candidate pool (47k models, size+family ids, popularity, HF urls)
 checkpoint/
-  MLPMetric.pt      ~37 MB trained weights
+  MLPMetricFull.pt  ~709 MB trained weights (slim: parent-class dead weights + train-set dataset_desc_matrix stripped)
   args.json         Training-time hyperparameters (model dims, num_*)
 data/
   task2id.json      Task vocab
   metric2id.json    Metric vocab
 ```
 
-The Space looks for the checkpoint at `checkpoint/MLPMetric.pt` and the data
-JSONs at `data/`. Override with env vars `MODEL_CKPT`, `MODEL_ARGS`, `DATA_DIR`,
-`POOL_PATH` if you lay things out differently.
+The Space looks for the checkpoint at `checkpoint/MLPMetricFull.pt` (or the
+legacy `checkpoint/MLPMetric.pt`) and the data JSONs at `data/`. Override with
+env vars `MODEL_CKPT`, `MODEL_ARGS`, `DATA_DIR`, `POOL_PATH` if you lay things
+out differently.
 
 ## Running locally
 
@@ -77,8 +80,12 @@ When you bump the candidate set (e.g. add new HF models to `model2id.json` /
 
 ```bash
 python web/build_model_pool.py \
-    --data-dir data/unified_augmented \
-    --args     checkpoint/mlp/unified_augmented/ablation_no_model_id_no_dataset_id/args.json \
-    --out      web/assets/model_pool.npz \
+    --data-dir    data/unified_augmented_v2 \
+    --profile-dir data/unified_augmented \
+    --args        checkpoint/mlp/unified_augmented_v2/FinalModel_v2_full_data_deployment/args.json \
+    --out         web/assets/model_pool.npz \
     --min-popularity 0
 ```
+
+(`--profile-dir` falls back to v1's `model_profile.json` / `model_popularity.json`
+for the ~21k v2 model names that v2 doesn't yet ship a profile for.)

app.py

@@ -140,7 +140,8 @@ with gr.Blocks(title="ModelLens · Finding the Best Model for Your Task", theme=
         # ModelLens: Finding the Best for Your Task from Myriads of Models
         Describe your dataset, pick a task type and a metric, and ModelLens returns
         the top candidates from a pool of **47k+** HuggingFace models. Backed by the
-        ablation_no_id MLPMetric checkpoint trained on `unified_augmented`.
+        `MLPMetricFull` checkpoint trained on the cleaned + expanded
+        `unified_augmented_v2` corpus.
 
         Results are post-filtered by a modality sanity check so that e.g.
         *Image Generation* won't surface text-only LLMs. The status line below

assets/model_pool.npz

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:66552520f9534fce6e4a530fe9ba55f8cf046d0c68ee0197eca02a988425c855
-size 5820984
+oid sha256:18f0046707d354836b60e244a54c4d84a9755b5ee984d875baff07f4c3185b14
+size 5802494

build_model_pool.py

@@ -62,6 +62,15 @@ def main(argv=None):
         default="checkpoint/mlp/unified_augmented/ablation_no_model_id_no_dataset_id/args.json",
         help="Path to the training args.json — used to read size_bucket so bucket ids align with the checkpoint.",
     )
+    p.add_argument(
+        "--profile-dir",
+        default=None,
+        help=(
+            "Optional fallback directory to read model_profile.json / "
+            "model_popularity.json from when --data-dir lacks them (e.g. "
+            "v2 deployment data only ships ID maps)."
+        ),
+    )
     p.add_argument("--out", default="web/assets/model_pool.npz")
     p.add_argument(
         "--min-popularity",
@@ -79,22 +88,36 @@ def main(argv=None):
         model2family = json.load(f)
     with open(os.path.join(args.data_dir, "family2id.json")) as f:
         family2id = json.load(f)
-    with open(os.path.join(args.data_dir, "model_profile.json")) as f:
-        model_profile = json.load(f)
-    pop_path = os.path.join(args.data_dir, "model_popularity.json")
-    pop_map = {}
-    if os.path.exists(pop_path):
-        pop_doc = json.load(open(pop_path))
-        # Doc shape: {fetched_at, source, num_models, status_counts, models: {name: {downloads, status}}}
-        models_field = pop_doc.get("models", pop_doc)
-        for name, entry in models_field.items():
-            if isinstance(entry, dict):
-                pop_map[name] = int(entry.get("downloads", 0) or 0)
-            else:
-                try:
-                    pop_map[name] = int(entry)
-                except Exception:
-                    pop_map[name] = 0
+
+    def _read_profile_files(d):
+        prof = {}
+        pop = {}
+        prof_path = os.path.join(d, "model_profile.json")
+        pop_path = os.path.join(d, "model_popularity.json")
+        if os.path.exists(prof_path):
+            with open(prof_path) as f:
+                prof = json.load(f)
+        if os.path.exists(pop_path):
+            pop_doc = json.load(open(pop_path))
+            models_field = pop_doc.get("models", pop_doc)
+            for name, entry in models_field.items():
+                if isinstance(entry, dict):
+                    pop[name] = int(entry.get("downloads", 0) or 0)
+                else:
+                    try:
+                        pop[name] = int(entry)
+                    except Exception:
+                        pop[name] = 0
+        return prof, pop
+
+    model_profile, pop_map = _read_profile_files(args.data_dir)
+    if args.profile_dir:
+        fb_prof, fb_pop = _read_profile_files(args.profile_dir)
+        # Fill in any gaps from the fallback dir (e.g. v1 profile for v2 names).
+        for k, v in fb_prof.items():
+            model_profile.setdefault(k, v)
+        for k, v in fb_pop.items():
+            pop_map.setdefault(k, v)
 
     if os.path.exists(args.args):
         train_args = json.load(open(args.args))

checkpoint/MLPMetricFull.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6c7a6ff547ee205e713593e3f0f539b6a646e8eaf02069c9fdfc8dfe052af9ee
+size 709051757

checkpoint/args.json

@@ -1 +1 @@
-{"device": "cuda:0", "use_data_parallel": false, "device_ids": [0, 1, 2, 3], "use_ddp": true, "ddp_find_unused_parameters": false, "num_workers": 0, "pin_memory": false, "persistent_workers": false, "data_name": "unified_augmented", "ood_split_mode": "new_dataset_evaluation", "seed": 2025, "use_wandb": true, "wandb_project": "ModelProfile", "wandb_entity": "ruicai-ucdavis", "trail_name": "ablation_no_model_id_no_dataset_id", "start_epoch": 0, "checkpoint_path": "", "is_train": true, "is_ood": true, "loss_type": "ensemble", "point_loss_weight": 0.1, "early_stop": 20, "num_epochs": 1000, "batch_size": 8, "pair_batch_size": 1024, "learning_rate": 0.001, "weight_decay": 0.0001, "tau": 10.0, "lambda_list": 0.5, "lambda_pair": 1.0, "alpha": 3.0, "size_bucket": [0.001, 0.003, 0.01, 0.03, 0.06, 0.1, 0.15, 0.2, 0.3, 0.4, 0.5, 0.6, 0.8, 1, 3, 7, 14, 35, 70, 100, 1000], "use_id_emb": false, "model_dim": 1536, "token_dim": 512, "use_size_prior": true, "size_dim": 64, "use_family_prior": true, "family_dim": 64, "dataset_desp_dim": 1536, "task_dim": 256, "model_name": "MLPMetric", "hidden_dim": 512, "dropout_rate": 0.02, "topk": [1, 3, 5, 7, 10, 30, 50, 70, 100], "margin_eps": 0.02, "val_eval_target_models_all_datasets": false, "val_eval_fixed_backbones": false, "save_best_ic8x10_checkpoint": false, "test_eval_target_models_all_datasets": false, "config": "config/ablations/MLPMetric_NoModelID_unified_augmented.yaml", "is_distributed": true, "world_size": 4, "rank": 0, "local_rank": 0, "num_models": 47062, "num_tasks": 2551, "num_metrics": 8420, "unknown_metric_id": 0, "num_size_buckets": 23, "num_families": 331}
+{"device": "cuda:0", "use_data_parallel": false, "device_ids": [0, 1, 2, 3], "use_ddp": true, "ddp_find_unused_parameters": false, "num_workers": 0, "pin_memory": false, "persistent_workers": false, "data_name": "unified_augmented_v2", "ood_split_mode": "new_dataset_evaluation", "test_split_mode": "val", "seed": 2025, "use_wandb": true, "wandb_project": "ModelProfile", "wandb_entity": "ruicai-ucdavis", "trail_name": "FinalModel_v2_full_data_deployment", "start_epoch": 0, "checkpoint_path": "", "is_train": true, "is_ood": false, "loss_type": "ensemble", "point_loss_weight": 0.1, "early_stop": 99999, "eval_every": 99999, "num_epochs": 30, "save_every": 5, "save_final_checkpoint": true, "batch_size": 8, "pair_batch_size": 1024, "learning_rate": 0.001, "weight_decay": 0.0001, "tau": 10.0, "lambda_list": 0.5, "lambda_pair": 1.0, "alpha": 3.0, "size_bucket": [0.001, 0.003, 0.01, 0.03, 0.06, 0.1, 0.15, 0.2, 0.3, 0.4, 0.5, 0.6, 0.8, 1, 3, 7, 14, 35, 70, 100, 1000], "use_id_emb": true, "model_dim": 1536, "token_dim": 512, "use_size_prior": true, "size_dim": 64, "use_family_prior": true, "family_dim": 64, "model_desp_emb_dim": 1536, "model_desp_emb_path": "data/unified_augmented_v2/model2desp_embeddings.npz", "use_dataset_id_as_desp": true, "dataset_desp_dim": 1, "dataset_id_emb_dim": 256, "dataset_desp_emb_dim": 1536, "task_dim": 256, "model_name": "MLPMetricFull", "hidden_dim": 512, "dropout_rate": 0.02, "id_dropout_rate": 0.1, "topk": [1, 10, 30, 50], "margin_eps": 0.02, "val_eval_target_models_all_datasets": false, "val_eval_fixed_backbones": false, "save_best_ic8x10_checkpoint": false, "test_eval_target_models_all_datasets": false, "config": "config/FinalModel_unified_augmented_v2.yaml", "is_distributed": true, "world_size": 4, "rank": 0, "local_rank": 0, "num_models": 47242, "num_tasks": 2581, "num_metrics": 3714, "num_datasets": 85937, "unknown_metric_id": 0, "num_size_buckets": 23, "num_families": 332}

data/family2id.json

@@ -272,62 +272,63 @@
   "singularity": 270,
   "sjt": 271,
   "slowfast": 272,
-  "smollm": 273,
-  "smoltulu": 274,
-  "solar": 275,
-  "sombrero": 276,
-  "speechstew": 277,
-  "stablelm": 278,
-  "starcoder": 279,
-  "stm": 280,
-  "summer": 281,
-  "svtr": 282,
-  "swin": 283,
-  "t5": 284,
-  "tarsier": 285,
-  "thea": 286,
-  "tinymistral": 287,
-  "tinyvit": 288,
-  "titannet": 289,
-  "tora": 290,
-  "transformer": 291,
-  "transnext": 292,
-  "triangulum": 293,
-  "trocr": 294,
-  "tsunami": 295,
-  "twist": 296,
-  "two": 297,
-  "ul2": 298,
-  "ultiima": 299,
-  "una": 300,
-  "unet": 301,
-  "unifiedqa": 302,
-  "uniformer": 303,
-  "uninet": 304,
-  "unireplknet": 305,
-  "uniter": 306,
-  "unknown": 307,
-  "van": 308,
-  "vgg": 309,
-  "vicious": 310,
-  "video": 311,
-  "vila": 312,
-  "vilt": 313,
-  "vinvl": 314,
-  "vit": 315,
-  "vlm": 316,
-  "wav2vec": 317,
-  "wav2vec2": 318,
-  "wavlm": 319,
-  "whisper": 320,
-  "wide": 321,
-  "winter": 322,
-  "wizard": 323,
-  "xcit": 324,
-  "xlm": 325,
-  "xlnet": 326,
-  "xmem": 327,
-  "yi": 328,
-  "zephyr": 329,
-  "zeus": 330
+  "slowfast,": 273,
+  "smollm": 274,
+  "smoltulu": 275,
+  "solar": 276,
+  "sombrero": 277,
+  "speechstew": 278,
+  "stablelm": 279,
+  "starcoder": 280,
+  "stm": 281,
+  "summer": 282,
+  "svtr": 283,
+  "swin": 284,
+  "t5": 285,
+  "tarsier": 286,
+  "thea": 287,
+  "tinymistral": 288,
+  "tinyvit": 289,
+  "titannet": 290,
+  "tora": 291,
+  "transformer": 292,
+  "transnext": 293,
+  "triangulum": 294,
+  "trocr": 295,
+  "tsunami": 296,
+  "twist": 297,
+  "two": 298,
+  "ul2": 299,
+  "ultiima": 300,
+  "una": 301,
+  "unet": 302,
+  "unifiedqa": 303,
+  "uniformer": 304,
+  "uninet": 305,
+  "unireplknet": 306,
+  "uniter": 307,
+  "unknown": 308,
+  "van": 309,
+  "vgg": 310,
+  "vicious": 311,
+  "video": 312,
+  "vila": 313,
+  "vilt": 314,
+  "vinvl": 315,
+  "vit": 316,
+  "vlm": 317,
+  "wav2vec": 318,
+  "wav2vec2": 319,
+  "wavlm": 320,
+  "whisper": 321,
+  "wide": 322,
+  "winter": 323,
+  "wizard": 324,
+  "xcit": 325,
+  "xlm": 326,
+  "xlnet": 327,
+  "xmem": 328,
+  "yi": 329,
+  "zephyr": 330,
+  "zeus": 331
 }

inference_lib.py

@@ -1,15 +1,17 @@
 """Self-contained inference module for the recommendation web app.
 
-Contains a trimmed copy of ``MLPMetric`` (and its dependencies) so HF Spaces
-deployments do not need to ship the full ``module/`` package. The class layout
-and parameter names match the trained checkpoint exactly, so the original
-``state_dict`` loads with ``strict=False`` and a clean diff.
+Contains trimmed copies of ``MLPMetric`` and ``MLPMetricFull`` (and their
+dependencies) so HF Spaces deployments do not need to ship the full
+``module/`` package. The class layout and parameter names match the trained
+checkpoint exactly, so the original ``state_dict`` loads with
+``strict=False`` and a clean diff.
 """
 from __future__ import annotations
 
 import hashlib
 import math
 import re
+from types import SimpleNamespace
 from typing import Optional
 
 import torch
@@ -248,3 +250,285 @@ class MLPMetric(nn.Module):
             out[:, start:end] = (s_chunk + prior_chunk) / T
             start = end
         return out
+
+
+class MLPMetricFull(MLPMetric):
+    """Full-feature recommender. Uses model-id emb, model-name emb, model-desc
+    emb, dataset-id emb, and dataset-desc emb.
+
+    For inference on a *new user dataset* (no global dataset_id), we:
+      - feed UNK as dataset_id (so dataset_id_embedding still contributes a
+        learned [UNK] prior),
+      - feed the user's OpenAI embedding directly as dataset_desc_emb,
+        bypassing the training-time ``dataset_desc_matrix`` lookup.
+
+    Parameter layout matches the training-time class so the state_dict loads
+    via ``load_state_dict(strict=False)`` after stripping the buffers that
+    are only useful for the train-set IDs.
+    """
+
+    def __init__(self, args):
+        # ---- dim bookkeeping ----
+        self.dataset_id_emb_dim = int(getattr(args, "dataset_id_emb_dim", 256))
+        self.dataset_desp_emb_dim = int(getattr(args, "dataset_desp_emb_dim", 1536))
+        self.model_desp_emb_dim = int(getattr(args, "model_desp_emb_dim", 1536))
+
+        # Information-source flags (kept for parity; defaults match training)
+        self.use_model_id_emb = bool(getattr(args, "use_model_id_emb", True))
+        self.use_model_name_emb = bool(getattr(args, "use_model_name_emb", True))
+        self.use_model_desc_emb = bool(getattr(args, "use_model_desc_emb", True))
+        self.use_dataset_id_emb = bool(getattr(args, "use_dataset_id_emb", True))
+        self.use_dataset_desc_emb = bool(getattr(args, "use_dataset_desc_emb", True))
+        self.use_size_feature = bool(getattr(args, "use_size_feature", True))
+
+        # The parent's __init__ builds task/size/family/metric embeddings,
+        # prior_head, temperature, plus a placeholder backbone (which we rebuild).
+        # Set dataset_desp_dim so parent sizes its placeholder correctly; we
+        # don't actually use the parent's backbone — we rebuild it below.
+        orig_desp_dim = args.dataset_desp_dim
+        args.dataset_desp_dim = self.dataset_id_emb_dim + self.dataset_desp_emb_dim
+        super().__init__(args)
+        args.dataset_desp_dim = orig_desp_dim
+
+        # ==== Model-side components (own name encoder + own id emb) ====
+        if self.use_model_name_emb:
+            args_name_only = SimpleNamespace(**vars(args))
+            args_name_only.use_id_emb = False
+            self._name_encoder = ModelNameAvgEncoder(args_name_only)
+        else:
+            self._name_encoder = None
+
+        if self.use_model_id_emb:
+            self._id_emb = nn.Embedding(args.num_models + 1, args.model_dim)
+            self.unk_model_id = args.num_models
+        else:
+            self._id_emb = None
+            self.unk_model_id = 0
+
+        # Model-description buffer: one row per known model.
+        if self.use_model_desc_emb:
+            self.register_buffer(
+                "model_desc_matrix",
+                torch.zeros(args.num_models, self.model_desp_emb_dim),
+            )
+        else:
+            self.register_buffer(
+                "model_desc_matrix",
+                torch.zeros(0, self.model_desp_emb_dim),
+            )
+
+        # ==== Dataset-side components ====
+        num_datasets = int(getattr(args, "num_datasets", 100000))
+        if self.use_dataset_id_emb:
+            # +2: one for [UNK], one for the upper slack (matches training)
+            self.dataset_id_embedding = nn.Embedding(num_datasets + 2, self.dataset_id_emb_dim)
+            self.unk_dataset_id = num_datasets + 1
+        else:
+            self.dataset_id_embedding = None
+            self.unk_dataset_id = 0
+
+        # ``dataset_desc_matrix`` is NOT registered at inference time — we use
+        # the user's OpenAI embedding directly. The stripped checkpoint also
+        # omits this buffer.
+
+        # ==== Recompute backbone input dim and rebuild ====
+        model_info_dim = (
+            (args.token_dim if self.use_model_name_emb else 0)
+            + (args.model_dim if self.use_model_id_emb else 0)
+            + (self.model_desp_emb_dim if self.use_model_desc_emb else 0)
+        )
+        self.model_info_dim = model_info_dim
+
+        dataset_emb_dim = (
+            (self.dataset_id_emb_dim if self.use_dataset_id_emb else 0)
+            + (self.dataset_desp_emb_dim if self.use_dataset_desc_emb else 0)
+        )
+        self.dataset_emb_dim = dataset_emb_dim
+        dataset_info_dim = dataset_emb_dim + args.task_dim
+        metric_dim = self.metric_dim if self.use_metric_embedding else 0
+        size_emb_dim_eff = args.size_dim if self.use_size_feature else 0
+        backbone_in = (
+            model_info_dim
+            + dataset_info_dim
+            + size_emb_dim_eff
+            + self.family_dim
+            + metric_dim
+        )
+        self.backbone = nn.Sequential(
+            nn.Linear(backbone_in, args.hidden_dim),
+            nn.ReLU(),
+            nn.Dropout(args.dropout_rate),
+            nn.Linear(args.hidden_dim, args.hidden_dim),
+            nn.ReLU(),
+            nn.Dropout(args.dropout_rate),
+        )
+
+        prior_in_actual = 0
+        if self.use_size_prior and self.use_size_feature:
+            prior_in_actual += args.size_dim
+        if self.use_family_prior:
+            prior_in_actual += self.family_dim
+        if prior_in_actual > 0:
+            self.prior_head = nn.Sequential(
+                nn.Linear(prior_in_actual, args.hidden_dim // 2),
+                nn.ReLU(),
+                nn.Linear(args.hidden_dim // 2, 1),
+            )
+
+    # ------------------------------------------------------------------
+    # Model-side encoding (used by build_model_cache)
+    # ------------------------------------------------------------------
+    def encode_model(
+        self, model_ids: torch.LongTensor, model_names: list[str],
+    ) -> torch.Tensor:
+        B = model_ids.shape[0]
+        device = model_ids.device
+        parts = []
+        if self.use_model_name_emb:
+            parts.append(self._name_encoder(model_ids, model_names))
+        if self.use_model_id_emb:
+            parts.append(self._id_emb(model_ids))
+        if self.use_model_desc_emb:
+            if self.model_desc_matrix.shape[0] > 0:
+                safe_ids = model_ids.clamp(0, self.model_desc_matrix.shape[0] - 1)
+                parts.append(self.model_desc_matrix[safe_ids])
+            else:
+                parts.append(torch.zeros(B, self.model_desp_emb_dim, device=device))
+        if not parts:
+            return torch.zeros(B, 0, device=device)
+        if len(parts) == 1:
+            return parts[0]
+        return torch.cat(parts, dim=-1)
+
+    @torch.no_grad()
+    def build_model_cache(
+        self,
+        all_model_names: list[str],
+        all_model_size_ids: torch.LongTensor,
+        all_model_family_ids: Optional[torch.LongTensor] = None,
+        device=None,
+    ):
+        if device is None:
+            device = next(self.parameters()).device
+        size_ids = all_model_size_ids.to(device=device, dtype=torch.long)
+        M = len(all_model_names)
+        assert size_ids.shape[0] == M
+        model_ids = torch.arange(M, device=device, dtype=torch.long)
+
+        h_model = self.encode_model(model_ids, all_model_names)
+        h_size = self.size_embedding(size_ids) if self.use_size_feature else None
+        cache = {"h_model": h_model, "h_size": h_size, "size_ids": size_ids}
+        if self.use_family_prior and all_model_family_ids is not None:
+            family_ids = all_model_family_ids.to(device=device, dtype=torch.long)
+            cache["h_family"] = self.family_embedding(family_ids)
+            cache["family_ids"] = family_ids
+        else:
+            cache["h_family"] = None
+            cache["family_ids"] = None
+        return cache
+
+    # ------------------------------------------------------------------
+    # Dataset-side encoding for inference: user's OpenAI emb + UNK id
+    # ------------------------------------------------------------------
+    def _encode_dataset_at_inference(
+        self, dataset_desp_emb: torch.Tensor,
+    ) -> torch.Tensor:
+        """``dataset_desp_emb`` is the user's OpenAI vector of shape
+        ``[B, dataset_desp_emb_dim]``. We pair it with a learned [UNK]
+        dataset-id embedding, matching the training-time concatenation order
+        (id_emb || desc_emb).
+        """
+        B = dataset_desp_emb.shape[0]
+        device = dataset_desp_emb.device
+        parts = []
+        if self.use_dataset_id_emb and self.dataset_id_embedding is not None:
+            unk = torch.full((B,), int(self.unk_dataset_id), dtype=torch.long, device=device)
+            parts.append(self.dataset_id_embedding(unk))
+        if self.use_dataset_desc_emb:
+            parts.append(dataset_desp_emb)
+        if not parts:
+            return torch.zeros(B, 0, device=device)
+        if len(parts) == 1:
+            return parts[0]
+        return torch.cat(parts, dim=-1)
+
+    # ------------------------------------------------------------------
+    # score_matrix at inference time
+    # ------------------------------------------------------------------
+    @torch.no_grad()
+    def score_matrix(
+        self,
+        task_ids: torch.LongTensor,
+        dataset_desp_batch: torch.Tensor,
+        model_cache: dict,
+        metric_ids: Optional[torch.LongTensor] = None,
+        chunk_size: int = 8192,
+    ) -> torch.Tensor:
+        """``dataset_desp_batch`` here is the OpenAI embedding ``[B, 1536]``."""
+        device = dataset_desp_batch.device
+        B = dataset_desp_batch.size(0)
+
+        h_task = self.task_embedding(task_ids)
+        h_data = self._encode_dataset_at_inference(dataset_desp_batch)
+        h_metric = self._metric_embed(metric_ids, B, device)
+
+        h_model_all = model_cache["h_model"]
+        h_size_all = model_cache["h_size"] if self.use_size_feature else None
+        h_family_all = model_cache.get("h_family")
+        M = h_model_all.size(0)
+
+        prior_parts_all = []
+        if self.use_size_prior and h_size_all is not None:
+            prior_parts_all.append(h_size_all)
+        if self.use_family_prior and h_family_all is not None:
+            prior_parts_all.append(h_family_all)
+        if prior_parts_all:
+            prior_inp_all = (
+                torch.cat(prior_parts_all, dim=-1) if len(prior_parts_all) > 1 else prior_parts_all[0]
+            )
+            prior_all = self.prior_head(prior_inp_all).squeeze(-1)
+        else:
+            prior_all = torch.zeros(M, device=device)
+
+        out = torch.empty(B, M, device=device)
+        T = torch.clamp(self.temperature, min=1e-3)
+
+        start = 0
+        while start < M:
+            end = min(start + chunk_size, M)
+            m = end - start
+
+            h_model = h_model_all[start:end]
+            h_model_exp = h_model.unsqueeze(0).expand(B, m, -1) if h_model.shape[1] > 0 else None
+            h_data_exp = h_data.unsqueeze(1).expand(B, m, -1) if h_data.shape[1] > 0 else None
+            h_task_exp = h_task.unsqueeze(1).expand(B, m, -1)
+            h_size_exp = (
+                h_size_all[start:end].unsqueeze(0).expand(B, m, -1)
+                if h_size_all is not None else None
+            )
+            h_metric_exp = (
+                h_metric.unsqueeze(1).expand(B, m, -1) if h_metric is not None else None
+            )
+
+            parts = []
+            if h_model_exp is not None:
+                parts.append(h_model_exp)
+            if h_data_exp is not None:
+                parts.append(h_data_exp)
+            if h_size_exp is not None:
+                parts.append(h_size_exp)
+            if h_family_all is not None:
+                h_family_exp = h_family_all[start:end].unsqueeze(0).expand(B, m, -1)
+                parts.append(h_family_exp)
+            parts.append(h_task_exp)
+            if h_metric_exp is not None:
+                parts.append(h_metric_exp)
+            residual_inp = torch.cat(parts, dim=-1)
+
+            h = self.backbone(residual_inp.reshape(B * m, -1))
+            s_chunk = self.pairwise_head(h).reshape(B, m)
+            prior_chunk = prior_all[start:end].unsqueeze(0)
+            out[:, start:end] = (s_chunk + prior_chunk) / T
+            start = end
+
+        return out

recommend.py

@@ -14,7 +14,7 @@ from typing import List, Optional
 import numpy as np
 import torch
 
-from inference_lib import MLPMetric
+from inference_lib import MLPMetric, MLPMetricFull
 
 
 EMBEDDING_MODEL = "text-embedding-3-small"  # Must match what was used during training.
@@ -509,14 +509,16 @@ class Recommender:
             dtype=np.int64,
         )
 
-        # Build the MLPMetric model with the same hyper-parameters used for training.
+        # Build the recommender model with the same hyper-parameters used for training.
         cfg = self._train_args
+        model_name = str(cfg.get("model_name", "MLPMetric"))
         model_args = SimpleNamespace(
             num_models=cfg.get("num_models", len(self.model_names)),
             num_tasks=cfg.get("num_tasks"),
             num_metrics=cfg.get("num_metrics"),
             num_size_buckets=cfg.get("num_size_buckets"),
             num_families=cfg.get("num_families"),
+            num_datasets=cfg.get("num_datasets", 100000),
             token_dim=cfg["token_dim"],
             model_dim=cfg["model_dim"],
             task_dim=cfg["task_dim"],
@@ -524,15 +526,28 @@ class Recommender:
             size_dim=cfg["size_dim"],
             family_dim=cfg.get("family_dim", cfg["size_dim"]),
             dataset_desp_dim=cfg["dataset_desp_dim"],
+            dataset_id_emb_dim=cfg.get("dataset_id_emb_dim", 256),
+            dataset_desp_emb_dim=cfg.get("dataset_desp_emb_dim", 1536),
+            model_desp_emb_dim=cfg.get("model_desp_emb_dim", 1536),
             hidden_dim=cfg["hidden_dim"],
             dropout_rate=cfg.get("dropout_rate", 0.0),
             use_id_emb=bool(cfg.get("use_id_emb", False)),
             use_size_prior=bool(cfg.get("use_size_prior", True)),
             use_family_prior=bool(cfg.get("use_family_prior", False)),
+            use_size_feature=bool(cfg.get("use_size_feature", True)),
             use_metric_feature=bool(cfg.get("use_metric_feature", True)),
+            use_model_id_emb=bool(cfg.get("use_model_id_emb", True)),
+            use_model_name_emb=bool(cfg.get("use_model_name_emb", True)),
+            use_model_desc_emb=bool(cfg.get("use_model_desc_emb", True)),
+            use_dataset_id_emb=bool(cfg.get("use_dataset_id_emb", True)),
+            use_dataset_desc_emb=bool(cfg.get("use_dataset_desc_emb", True)),
             unknown_metric_id=int(cfg.get("unknown_metric_id", 0)),
         )
-        self.model = MLPMetric(model_args).to(self.device).eval()
+        if model_name == "MLPMetricFull":
+            self.model = MLPMetricFull(model_args).to(self.device).eval()
+        else:
+            self.model = MLPMetric(model_args).to(self.device).eval()
+        self._model_name = model_name
 
         raw = torch.load(checkpoint_path, map_location="cpu")
         state = raw.get("model", raw) if isinstance(raw, dict) else raw
@@ -766,13 +781,16 @@ def default_recommender() -> Recommender:
     here = os.path.dirname(os.path.abspath(__file__))
     root = os.path.dirname(here)
 
-    spaces_ckpt = os.path.join(here, "checkpoint/MLPMetric.pt")
+    # Prefer the v2 MLPMetricFull checkpoint name; fall back to legacy MLPMetric.pt.
     spaces_args = os.path.join(here, "checkpoint/args.json")
     spaces_data = os.path.join(here, "data")
+    spaces_ckpt_full = os.path.join(here, "checkpoint/MLPMetricFull.pt")
+    spaces_ckpt_metric = os.path.join(here, "checkpoint/MLPMetric.pt")
+    spaces_ckpt = spaces_ckpt_full if os.path.exists(spaces_ckpt_full) else spaces_ckpt_metric
 
-    dev_ckpt = os.path.join(root, "checkpoint/mlp/unified_augmented/ablation_no_model_id_no_dataset_id/MLPMetric.pt")
-    dev_args = os.path.join(root, "checkpoint/mlp/unified_augmented/ablation_no_model_id_no_dataset_id/args.json")
-    dev_data = os.path.join(root, "data/unified_augmented")
+    dev_ckpt = os.path.join(root, "checkpoint/mlp/unified_augmented_v2/FinalModel_v2_full_data_deployment/MLPMetricFull.pt")
+    dev_args = os.path.join(root, "checkpoint/mlp/unified_augmented_v2/FinalModel_v2_full_data_deployment/args.json")
+    dev_data = os.path.join(root, "data/unified_augmented_v2")
 
     def _pick(env_key: str, primary: str, fallback: str) -> str:
         v = os.environ.get(env_key)