From 5c79e63dd2d8b8fcaac05dd60464f166d70c0eb5 Mon Sep 17 00:00:00 2001 From: Peng Ding Date: Tue, 29 Sep 2026 20:54:10 -0500 Subject: [PATCH 1/2] Route augmented noul items through choice head for training Phase C of #104: training integration for augmented noul questions. - Add _select_teacher_probs() to data/pipeline.py: selects best teacher from nested multi-teacher cache (jev_augmented > gpt_5_6_luna_augmented for augmented noul, jev > luna > qwen for standard items). Fixes existing dead code where synthetic teacher_probs were never used. - Augmented noul items with 5-option teacher_probs are routed through the choice head (KL-divergence against teacher distribution) instead of the binary noul head, in both training and eval paths. - Uncertainty weighting correctly uses "choice" log_var for augmented noul items. - Eval maps augmented noul predictions back to binary accuracy (pred_key == "true"/"false"). - Standard noul fallback when teacher_probs has 2-key {"true","false"} format (uses true probability as BCE target). - 9 new tests for teacher selection logic, 557 total pass. --- data/pipeline.py | 32 ++++++- model/training/supervised.py | 117 ++++++++++++++++++++++---- tests/test_augmented_noul_training.py | 87 +++++++++++++++++++ 3 files changed, 217 insertions(+), 19 deletions(-) create mode 100644 tests/test_augmented_noul_training.py diff --git a/data/pipeline.py b/data/pipeline.py index 4ab340a..6c1996a 100644 --- a/data/pipeline.py +++ b/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) diff --git a/model/training/supervised.py b/model/training/supervised.py index e771ed6..bcdb15f 100644 --- a/model/training/supervised.py +++ b/model/training/supervised.py @@ -98,15 +98,26 @@ 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) + 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 +155,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 +208,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 +227,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 +237,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 +363,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 +576,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"] diff --git a/tests/test_augmented_noul_training.py b/tests/test_augmented_noul_training.py new file mode 100644 index 0000000..f5c4b47 --- /dev/null +++ b/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"] From 320ca4de36fb1ed80510f7a3623dff89bf463f59 Mon Sep 17 00:00:00 2001 From: Peng Ding Date: Thu, 8 Oct 2026 20:47:48 -0500 Subject: [PATCH 2/2] Add comment on q_type reassignment for uncertainty weighting MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses non-blocking review feedback from clementine and elena: the q_type = "choice" override in compute_loss is subtle — clarify that it ensures augmented noul uses the "choice" log_var. --- model/training/supervised.py | 1 + 1 file changed, 1 insertion(+) diff --git a/model/training/supervised.py b/model/training/supervised.py index bcdb15f..c0110d0 100644 --- a/model/training/supervised.py +++ b/model/training/supervised.py @@ -104,6 +104,7 @@ def compute_loss( 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: logit = model.forward_noul(state, instructions)