From 2c40f2966d67cacaf7308fe099dcd5e087f7d676 Mon Sep 17 00:00:00 2001 From: Yizheng Huang Date: Tue, 15 Sep 2026 16:19:24 -0700 Subject: [PATCH] generator: correct failure feedback and paper semantics --- proeval/generator/core.py | 65 +++++++++++++++++++++++-------------- proeval/generator/prompt.py | 12 +++---- 2 files changed, 46 insertions(+), 31 deletions(-) diff --git a/proeval/generator/core.py b/proeval/generator/core.py index 0d8a6ee..331ff01 100644 --- a/proeval/generator/core.py +++ b/proeval/generator/core.py @@ -179,9 +179,11 @@ def select_hard_problems_bq( threshold: float = 0.7, noise_variance: float = 0.3, ) -> Tuple[List[int], np.ndarray, np.ndarray]: - """Select hard problems via HSS active sampling. + """Select hard problems via SS active sampling on failure scores. - Returns ``(hard_indices, posterior_mean, posterior_var)``. + ``threshold`` is the failure-score boundary used for acquisition and + filtering the final posterior means. Returns + ``(hard_indices, posterior_mean, posterior_var)``. """ from proeval.sampler.bq import _get_posterior @@ -193,7 +195,7 @@ def select_hard_problems_bq( for _ in range(min(budget, n_samples)): if not unlabeled: break - best = ss_acquisition(u_t, s_t, unlabeled, threshold=0.5, beta=1.96) + best = ss_acquisition(u_t, s_t, unlabeled, threshold=threshold, beta=1.96) labeled.append(best) unlabeled.remove(best) if labeled: @@ -201,9 +203,9 @@ def select_hard_problems_bq( test_x[:, labeled], test_y[labeled], test_x, noise_variance, labeled, u ) - hard = [i for i in labeled if u[i] < threshold] + hard = [i for i in labeled if u_t[i] >= threshold] if len(hard) < budget // 2: - hard = sorted(labeled, key=lambda x: u[x])[:budget] + hard = sorted(labeled, key=lambda x: u_t[x], reverse=True)[:budget] return hard, u_t, s_t @@ -343,9 +345,9 @@ def get_posterior_embedding( phi_t = torch.from_numpy(phi_test).float().to(device) K = compute_kernel_matrix(phi_t, encoder) if full_cov: - return u.copy(), K.cpu().numpy() + return u.copy(), K.detach().cpu().numpy() else: - return u.copy(), torch.diag(K).cpu().numpy() + return u.copy(), torch.diag(K).detach().cpu().numpy() phi_train_t = torch.from_numpy(phi_train).float().to(device) phi_test_t = torch.from_numpy(phi_test).float().to(device) @@ -415,9 +417,11 @@ def get_posterior_embedding( class TopicAwareGenerator: """Topic-aware test case generator with internal GP state management. - Manages its own GP posterior and dynamically selects hard anchors via - SS acquisition each ``generate()`` call. Call ``update(score)`` after - evaluating each generated case to update the GP posterior. + Selects hard anchors via SS acquisition each ``generate()`` call. + Call ``update(score)`` after evaluating each generated case to update + topic rewards and score history. Generated-input GP conditioning is + not supported by this score-only API: it would also require embeddings + or kernel covariances for the generated inputs. Topic modeling (BERTopic) is handled internally — just pass ``n_topics``. @@ -513,6 +517,7 @@ def __init__( self.labeled_indices: List[int] = [] self.labeled_y: List[float] = [] self._last_anchors: List[int] = [] + self._last_topic: Optional[str] = None self._iteration = 0 # Encoder state (with pretrain only) @@ -608,6 +613,10 @@ def _select_topic_ucb1( where rĖ„(s) is the failure rate for topic s, N is total trials, and n(s) is trials for topic s. """ + for item in items: + if stats[item]["total"] == 0: + return item + N = max(1, sum(s["total"] for s in stats.values())) scores = [] for item in items: @@ -661,16 +670,22 @@ def _select_hard_anchors( # GP Posterior Update def update(self, score: float) -> None: - """Feed back an evaluation result to update the GP posterior. + """Record the latest generated case's score and topic reward. + + Generated inputs do not belong to the source pool, so their scores + alone cannot condition its GP. This method does not add generated + observations to the GP without their embeddings or covariances. Args: score: Error score — ``1.0`` for failure, ``0.0`` for correct. + Scores at or above ``0.5`` count as failures, matching + :attr:`failures_found`. """ self.labeled_y.append(score) - # Update topic stats for last generation - # (topic_stats already updated in generate(), but labeled_y is needed - # for posterior update) + if self._last_topic is not None: + self.update_stats(self._last_topic, score) + self._last_topic = None if self.prior_mode == "encoder": # With pretrain: encoder-based posterior (TPF) @@ -746,16 +761,10 @@ def generate(self, strategy: str = "tss", k_examples: int = 5) -> Dict: # Topic selection if strategy in ("pure_random", "ss_gen"): topic = None - topic_id = None elif strategy in ("random_topic", "random"): topic = random.choice(self.topics) - topic_id = self.topics.index(topic) if topic in self.topics else None - if topic_id is not None and topic_id < len(self.unique_topics): - topic_id = self.unique_topics[topic_id % len(self.unique_topics)] elif strategy == "tss": topic = self._select_topic_ucb1(self.topics, self.topic_stats) - topic_idx = self.topics.index(topic) if topic in self.topics else 0 - topic_id = self.unique_topics[topic_idx % len(self.unique_topics)] else: raise ValueError( f"Unknown strategy: {strategy!r}. " @@ -764,9 +773,10 @@ def generate(self, strategy: str = "tss", k_examples: int = 5) -> Dict: # Hard anchor selection if strategy in ("tss", "ss_gen") and k_examples > 0: - use_topic = strategy == "tss" + # TSS transfers failure patterns from any topic into the chosen + # topic; anchor selection is independent of topic selection. selected_hard = self._select_hard_anchors( - k_examples, topic_id=topic_id, use_topic=use_topic, + k_examples, use_topic=False, ) else: selected_hard = [] @@ -810,14 +820,19 @@ def generate(self, strategy: str = "tss", k_examples: int = 5) -> Dict: # Update topic stats if topic and topic in self.topic_stats: self.topic_stats[topic]["total"] += 1 + self._last_topic = topic return result def update_stats(self, topic: str, score: float) -> None: """Record a failure for a topic (called after evaluation). - .. deprecated:: Use :meth:`update` instead, which handles both - topic stats and GP posterior. + ``score`` follows the error convention: ``1.0`` for failure and + ``0.0`` for correct. Values at or above ``0.5`` count as failures, + matching :attr:`failures_found`. + + .. deprecated:: Use :meth:`update` instead to also record score + history for the latest generated case. """ - if topic in self.topic_stats and score == 0.0: + if topic in self.topic_stats and score >= 0.5: self.topic_stats[topic]["failures"] += 1 diff --git a/proeval/generator/prompt.py b/proeval/generator/prompt.py index 1f64104..40993d4 100644 --- a/proeval/generator/prompt.py +++ b/proeval/generator/prompt.py @@ -82,9 +82,9 @@ def format_hard_examples(hard_examples: List[Dict]) -> Tuple[str, float]: for i, ex in enumerate(hard_examples): pm = ex.get("prior_mean", 0.5) avg += pm - diff = "VERY HARD" if pm < 0.3 else "HARD" if pm < 0.5 else "MODERATE" + diff = "VERY HARD" if pm >= 0.7 else "HARD" if pm >= 0.5 else "MODERATE" ex_text += ( - f"--- Example {i + 1} [{diff}, success rate: {pm:.1%}] ---\n" + f"--- Example {i + 1} [{diff}, predicted failure rate: {pm:.1%}] ---\n" f"Question: {ex['question']}\nAnswer: {ex['ground_truth']}\n\n" ) if hard_examples: @@ -120,7 +120,7 @@ def build_gsm8k_prompt(topic, hard_examples, strategy) -> str: return ( "You are an expert Red-Teamer creating math problems AI models get WRONG.\n\n" f"=== HARD EXAMPLES ===\n{ex_text}" - f"Average model success rate: {avg:.1%}\n\n" + f"Average predicted failure rate: {avg:.1%}\n\n" "Generate ONE NEW math problem EQUALLY DIFFICULT or HARDER.\n\n" "Requirements:\n" "1. Mimic the reasoning pattern of the hard examples\n" @@ -135,7 +135,7 @@ def build_gsm8k_prompt(topic, hard_examples, strategy) -> str: return ( "You are an expert Red-Teamer creating math problems AI models get WRONG.\n\n" f"=== HARD EXAMPLES ===\n{ex_text}" - f"Average model success rate: {avg:.1%}\n\n" + f"Average predicted failure rate: {avg:.1%}\n\n" f"Generate ONE NEW math problem EQUALLY DIFFICULT or HARDER.\nTOPIC: {topic}\n\n" "Requirements:\n" "1. Mimic the reasoning pattern of the hard examples\n" @@ -181,7 +181,7 @@ def build_strategyqa_prompt(topic, hard_examples, strategy) -> str: return ( "You are an expert Red-Teamer creating yes/no questions AI models get WRONG.\n\n" f"=== HARD EXAMPLES (AI models failed on these) ===\n{ex_text}" - f"Average model success rate: {avg:.1%}\n\n" + f"Average predicted failure rate: {avg:.1%}\n\n" "Generate ONE NEW yes/no question EQUALLY DIFFICULT or HARDER.\n\n" "Requirements:\n" "1. Mimic the reasoning pattern of the hard examples\n" @@ -196,7 +196,7 @@ def build_strategyqa_prompt(topic, hard_examples, strategy) -> str: return ( "You are an expert Red-Teamer creating yes/no questions AI models get WRONG.\n\n" f"=== HARD EXAMPLES (AI models failed on these) ===\n{ex_text}" - f"Average model success rate: {avg:.1%}\n\n" + f"Average predicted failure rate: {avg:.1%}\n\n" f"Generate ONE NEW yes/no question EQUALLY DIFFICULT or HARDER.\n" f"TOPIC: {topic}\n\n" "Requirements:\n"