Route augmented noul through choice head for training (Phase C)

data/pipeline.py

@@ -1028,6 +1028,34 @@ def load_drop() -> Iterator[TypedQuestion]:
 # --- Synthetic (generated data) ---
 
 
+_AUGMENTED_TEACHER_PRIORITY = ("jev_augmented", "gpt_5_6_luna_augmented")
+_STANDARD_TEACHER_PRIORITY = ("jev", "gpt_5_6_luna", "qwen3_8_27b")
+
+
+def _select_teacher_probs(
+    cached: dict[str, dict[str, float]],
+    question: dict[str, Any],
+) -> dict[str, float] | None:
+    """Select the best teacher probs from the nested multi-teacher cache.
+
+    For augmented noul items (question has augmented_options), prefer
+    augmented teachers. For standard items, prefer Jev > Luna > Qwen.
+    Returns a flat ``{option_key: probability}`` dict or None.
+    """
+    is_augmented_noul = (
+        question.get("type") == "noul"
+        and question.get("augmented_options")
+    )
+    priority = _AUGMENTED_TEACHER_PRIORITY if is_augmented_noul else _STANDARD_TEACHER_PRIORITY
+    for key in priority:
+        if key in cached and cached[key]:
+            return cached[key]
+    for key in _STANDARD_TEACHER_PRIORITY:
+        if key in cached and cached[key]:
+            return cached[key]
+    return None
+
+
 def load_synthetic() -> Iterator[TypedQuestion]:
     """Load synthetic data by assembling from families + variants + labels on the fly."""
     synth_dir = DATA_DIR / "synthetic"
@@ -1058,7 +1086,9 @@ def load_synthetic() -> Iterator[TypedQuestion]:
             for item_dict in items:
                 cached = label_cache.get(item_dict["id"])
                 if cached:
-                    item_dict["teacher_probs"] = cached
+                    flat = _select_teacher_probs(cached, item_dict.get("question", {}))
+                    if flat:
+                        item_dict["teacher_probs"] = flat
                 yield TypedQuestion(**item_dict)
 
 

model/training/supervised.py

@@ -98,15 +98,27 @@ def compute_loss(
     loss: torch.Tensor | None = None
 
     if q_type == "noul":
-        logit = model.forward_noul(state, instructions)
-        if isinstance(label, str):
-            label = label.lower() in ("true", "yes", "1")
-        if teacher_probs and "noul" in teacher_probs:
-            teacher_p = float(teacher_probs["noul"])
-            target = torch.tensor([[teacher_p]], device=logit.device)
+        augmented_options = question.get("augmented_options")
+        if augmented_options and teacher_probs and len(teacher_probs) > 2:
+            keys = list(augmented_options.keys())
+            option_texts = [augmented_options[k] for k in keys]
+            logits = model.forward_choice(state, instructions, option_texts)
+            loss = _teacher_loss(logits, teacher_probs, keys)
+            # Override q_type so uncertainty weighting uses "choice" log_var
+            q_type = "choice"
         else:
-            target = torch.tensor([[1.0 if label else 0.0]], device=logit.device)
-        loss = F.binary_cross_entropy_with_logits(logit, target)
+            logit = model.forward_noul(state, instructions)
+            if isinstance(label, str):
+                label = label.lower() in ("true", "yes", "1")
+            if teacher_probs and "noul" in teacher_probs:
+                teacher_p = float(teacher_probs["noul"])
+                target = torch.tensor([[teacher_p]], device=logit.device)
+            elif teacher_probs and len(teacher_probs) == 2 and "true" in teacher_probs:
+                teacher_p = float(teacher_probs["true"])
+                target = torch.tensor([[teacher_p]], device=logit.device)
+            else:
+                target = torch.tensor([[1.0 if label else 0.0]], device=logit.device)
+            loss = F.binary_cross_entropy_with_logits(logit, target)
 
     elif q_type == "choice":
         criteria = question["criteria"]
@@ -144,18 +156,49 @@ def compute_loss(
     return loss
 
 
+def _is_augmented_noul(item: Any) -> bool:
+    """Check if a noul item has augmented options with multi-option teacher probs."""
+    ao = item.question.get("augmented_options")
+    tp = getattr(item, "teacher_probs", None)
+    return bool(ao and tp and len(tp) > 2)
+
+
 def _compute_noul_batch(
     model: nn.Module,
     items: list,
 ) -> list[torch.Tensor]:
     """Batch-compute noul losses: one backbone call for all contexts.
 
+    Augmented noul items (with augmented_options + multi-option teacher_probs)
+    are routed through the choice head instead of the binary noul head.
+
     When the model uses a unified head (model.noul_head is None), falls back
     to per-item forward_noul calls since noul requires cross-attention with
     "true"/"false" option embeddings rather than a simple linear head.
     """
+    standard = []
+    augmented = []
+    item_order: list[tuple[str, int]] = []
+    for item in items:
+        if _is_augmented_noul(item):
+            item_order.append(("aug", len(augmented)))
+            augmented.append(item)
+        else:
+            item_order.append(("std", len(standard)))
+            standard.append(item)
+
+    std_losses = _compute_standard_noul_batch(model, standard) if standard else []
+    aug_losses = _compute_augmented_noul_batch(model, augmented) if augmented else []
+
+    return [std_losses[idx] if tag == "std" else aug_losses[idx] for tag, idx in item_order]
+
+
+def _compute_standard_noul_batch(
+    model: nn.Module,
+    items: list,
+) -> list[torch.Tensor]:
+    """Binary noul path: original batch logic."""
     if getattr(model, "noul_head", None) is None:
-        # Unified head path: per-item forward through cross-attention
         losses: list[torch.Tensor] = []
         for item in items:
             instructions = item.question.get("instructions", "")
@@ -166,12 +209,13 @@ def _compute_noul_batch(
             teacher_probs = getattr(item, "teacher_probs", None)
             if teacher_probs and "noul" in teacher_probs:
                 target = torch.tensor([[float(teacher_probs["noul"])]], device=logit.device)
+            elif teacher_probs and len(teacher_probs) == 2 and "true" in teacher_probs:
+                target = torch.tensor([[float(teacher_probs["true"])]], device=logit.device)
             else:
                 target = torch.tensor([[1.0 if label else 0.0]], device=logit.device)
             losses.append(F.binary_cross_entropy_with_logits(logit, target))
         return losses
 
-    # Standard batch path: one backbone call for all noul contexts
     texts = [f"{it.state} {it.question.get('instructions', '')}" for it in items]
     pooled = model._encode_text(texts)  # [N, hidden]
     logits = model.noul_head(pooled)  # [N, 1]
@@ -184,6 +228,8 @@ def _compute_noul_batch(
         teacher_probs = getattr(item, "teacher_probs", None)
         if teacher_probs and "noul" in teacher_probs:
             targets.append(float(teacher_probs["noul"]))
+        elif teacher_probs and len(teacher_probs) == 2 and "true" in teacher_probs:
+            targets.append(float(teacher_probs["true"]))
         else:
             targets.append(1.0 if label else 0.0)
 
@@ -192,6 +238,28 @@ def _compute_noul_batch(
     return [batch_loss[i] for i in range(len(items))]
 
 
+def _compute_augmented_noul_batch(
+    model: nn.Module,
+    items: list,
+) -> list[torch.Tensor]:
+    """Route augmented noul items through the choice head with 5-option softmax."""
+    losses: list[torch.Tensor] = []
+    for item in items:
+        ao = item.question["augmented_options"]
+        keys = list(ao.keys())
+        option_texts = [ao[k] for k in keys]
+        teacher_probs = item.teacher_probs
+
+        context_text = f"{item.state} {item.question.get('instructions', '')}"
+        context_hidden = model._encode_with_sequence(context_text)
+        option_pooled = model._encode_text(option_texts)
+        opt_hidden = option_pooled.unsqueeze(0)
+        logits = model.choice_head(context_hidden, opt_hidden)
+
+        losses.append(_teacher_loss(logits, teacher_probs, keys))
+    return losses
+
+
 def _compute_choice_batch(
     model: nn.Module,
     items: list,
@@ -296,10 +364,12 @@ def train_epoch(
 
             if noul_items:
                 noul_losses = _compute_noul_batch(model, noul_items)
-                for loss in noul_losses:
+                for i_noul, loss in enumerate(noul_losses):
                     if loss is not None:
-                        if log_vars is not None and "noul" in log_vars:
-                            loss = _apply_uncertainty_weight(loss, log_vars["noul"])
+                        if log_vars is not None:
+                            lv_key = "choice" if _is_augmented_noul(noul_items[i_noul]) else "noul"
+                            if lv_key in log_vars:
+                                loss = _apply_uncertainty_weight(loss, log_vars[lv_key])
                         all_weighted.append(loss.squeeze())
                     else:
                         n_window_skipped += 1
@@ -507,11 +577,23 @@ def _eval_item(
     if q_type == "noul":
         if isinstance(label, str):
             label = label.lower() in ("true", "yes", "1")
-        logit = model.forward_noul(state, instructions)
-        target = torch.tensor([[1.0 if label else 0.0]], device=logit.device)
-        loss = F.binary_cross_entropy_with_logits(logit, target)
-        pred = torch.sigmoid(logit).item() > 0.5
-        return (loss.item(), pred == label)
+        augmented_options = question.get("augmented_options")
+        teacher_probs = getattr(item, "teacher_probs", None)
+        if augmented_options and teacher_probs and len(teacher_probs) > 2:
+            keys = list(augmented_options.keys())
+            option_texts = [augmented_options[k] for k in keys]
+            logits = model.forward_choice(state, instructions, option_texts)
+            loss = _teacher_loss(logits, teacher_probs, keys)
+            pred_idx = logits.argmax(dim=-1).item()
+            pred_key = keys[pred_idx]
+            target_key = "true" if label else "false"
+            return (loss.item(), pred_key == target_key)
+        else:
+            logit = model.forward_noul(state, instructions)
+            target = torch.tensor([[1.0 if label else 0.0]], device=logit.device)
+            loss = F.binary_cross_entropy_with_logits(logit, target)
+            pred = torch.sigmoid(logit).item() > 0.5
+            return (loss.item(), pred == label)
 
     elif q_type == "choice":
         criteria = question["criteria"]

tests/test_augmented_noul_training.py

@@ -0,0 +1,87 @@
+"""Tests for Phase C: augmented noul training integration.
+
+Covers teacher selection/flattening in pipeline.py and augmented noul
+routing through the choice head in supervised.py.
+"""
+
+from __future__ import annotations
+
+import pytest
+
+from data.pipeline import _select_teacher_probs
+
+
+# ---------------------------------------------------------------------------
+# _select_teacher_probs tests
+# ---------------------------------------------------------------------------
+
+
+class TestSelectTeacherProbs:
+    def test_augmented_noul_prefers_jev_augmented(self):
+        cached = {
+            "jev": {"true": 0.8, "false": 0.2},
+            "jev_augmented": {"true": 0.6, "false": 0.1, "partial": 0.2, "unlikely": 0.05, "unrelated": 0.05},
+            "gpt_5_6_luna_augmented": {"true": 0.5, "false": 0.2, "partial": 0.15, "unlikely": 0.1, "unrelated": 0.05},
+        }
+        question = {"type": "noul", "instructions": "test?", "augmented_options": {"true": "T", "false": "F", "partial": "P", "unlikely": "U", "unrelated": "X"}}
+        result = _select_teacher_probs(cached, question)
+        assert result == cached["jev_augmented"]
+
+    def test_augmented_noul_falls_back_to_luna_augmented(self):
+        cached = {
+            "jev": {"true": 0.8, "false": 0.2},
+            "gpt_5_6_luna_augmented": {"true": 0.5, "false": 0.2, "partial": 0.15, "unlikely": 0.1, "unrelated": 0.05},
+        }
+        question = {"type": "noul", "instructions": "test?", "augmented_options": {"true": "T", "false": "F", "partial": "P", "unlikely": "U", "unrelated": "X"}}
+        result = _select_teacher_probs(cached, question)
+        assert result == cached["gpt_5_6_luna_augmented"]
+
+    def test_augmented_noul_falls_back_to_standard_if_no_augmented_teacher(self):
+        cached = {
+            "jev": {"true": 0.8, "false": 0.2},
+        }
+        question = {"type": "noul", "instructions": "test?", "augmented_options": {"true": "T", "false": "F", "partial": "P"}}
+        result = _select_teacher_probs(cached, question)
+        assert result == cached["jev"]
+
+    def test_standard_noul_prefers_jev(self):
+        cached = {
+            "jev": {"true": 0.8, "false": 0.2},
+            "gpt_5_6_luna": {"true": 0.7, "false": 0.3},
+        }
+        question = {"type": "noul", "instructions": "test?"}
+        result = _select_teacher_probs(cached, question)
+        assert result == cached["jev"]
+
+    def test_standard_noul_falls_back_to_luna(self):
+        cached = {
+            "gpt_5_6_luna": {"true": 0.7, "false": 0.3},
+        }
+        question = {"type": "noul", "instructions": "test?"}
+        result = _select_teacher_probs(cached, question)
+        assert result == cached["gpt_5_6_luna"]
+
+    def test_choice_prefers_jev(self):
+        cached = {
+            "jev": {"a": 0.5, "b": 0.3, "c": 0.2},
+            "gpt_5_6_luna": {"a": 0.4, "b": 0.4, "c": 0.2},
+        }
+        question = {"type": "choice", "instructions": "pick one", "criteria": {"a": "A", "b": "B", "c": "C"}}
+        result = _select_teacher_probs(cached, question)
+        assert result == cached["jev"]
+
+    def test_empty_cache_returns_none(self):
+        result = _select_teacher_probs({}, {"type": "noul", "instructions": "x?"})
+        assert result is None
+
+    def test_empty_teacher_dict_skipped(self):
+        cached = {"jev": {}, "gpt_5_6_luna": {"true": 0.6, "false": 0.4}}
+        question = {"type": "noul", "instructions": "test?"}
+        result = _select_teacher_probs(cached, question)
+        assert result == cached["gpt_5_6_luna"]
+
+    def test_qwen_fallback(self):
+        cached = {"qwen3_8_27b": {"true": 0.9, "false": 0.1}}
+        question = {"type": "noul", "instructions": "test?"}
+        result = _select_teacher_probs(cached, question)
+        assert result == cached["qwen3_8_27b"]