| 253 | ] |
| 254 | |
| 255 | def predict(self, row: dict[str, Any]) -> TierVote: |
| 256 | if self._embed_fn is None: |
| 257 | return TierVote(tier_id=None, confidence=0.0) |
| 258 | |
| 259 | messages = row.get("messages", []) |
| 260 | text = _extract_last_user_message(messages) |
| 261 | if not text.strip(): |
| 262 | return TierVote(tier_id=None, confidence=0.0) |
| 263 | |
| 264 | query_vec = self._embed_fn(text) |
| 265 | meta_feats = self._extract_meta_features(messages, text) |
| 266 | |
| 267 | # If the classifier was trained on dual embeddings (v0 = last_user, |
| 268 | # v1 = agent_state), compute the second embedding too. We detect this |
| 269 | # by the classifier's n_features_in_: 399 = single (384 + 15), 783 = dual. |
| 270 | state_vec = None |
| 271 | if self._classifier is not None and hasattr(self._classifier, "n_features_in_"): |
| 272 | expected = int(self._classifier.n_features_in_) |
| 273 | single_dim = int(query_vec.shape[-1]) + len(meta_feats) |
| 274 | if expected == single_dim + int(query_vec.shape[-1]): |
| 275 | state_text = _extract_agent_state(messages) |
| 276 | if state_text: |
| 277 | state_vec = self._embed_fn(state_text) |
| 278 | else: |
| 279 | # Fall back to duplicating the user embedding so dim matches. |
| 280 | state_vec = query_vec |
| 281 | |
| 282 | # Hybrid: classifier first, KNN fallback when classifier is uncertain. |
| 283 | if self._classifier is not None: |
| 284 | vote = self._predict_classifier(query_vec, meta_feats, state_vec=state_vec) |
| 285 | if vote.confidence >= self._clf_fallback_threshold: |
| 286 | return vote |
| 287 | if self._embeddings is not None and self._labels is not None: |
| 288 | knn_vote = self._predict_knn(query_vec) |
| 289 | # When kNN abstains (e.g., embedding collision with mixed labels), |
| 290 | # trust the classifier's meta-feature-aware prediction rather |
| 291 | # than returning None. This matters for multi-step swebench rows |
| 292 | # where the last_user text is identical across steps but the |
| 293 | # per-step difficulty varies. |
| 294 | if knn_vote.tier_id is None and vote.tier_id is not None: |
| 295 | return vote |
| 296 | return knn_vote |
| 297 | return vote |
| 298 | |
| 299 | # No classifier — KNN only |
| 300 | if self._embeddings is None or self._labels is None: |
| 301 | return TierVote(tier_id=None, confidence=0.0) |
| 302 | return self._predict_knn(query_vec) |
| 303 | |
| 304 | def _predict_classifier(self, query_vec: np.ndarray, meta_feats: list[float] | None = None, |
| 305 | state_vec: np.ndarray | None = None) -> TierVote: |