Merge: F2 backend — emit effective_probability for sampling-effective margins (#5)

backend/model_service.py

@@ -2176,6 +2176,21 @@ async def analyze_research_attention(request: Dict[str, Any], authenticated: boo
                 log_probs = torch.nn.functional.log_softmax(logits, dim=-1)
                 top_probs = torch.exp(log_probs[top_indices])
 
+                # `probability` is the post-temperature, PRE-filter probability
+                # (model preference). `effective_probability` is the post-filter
+                # renormalised probability — what the sampler actually drew
+                # from. Under top-k / top-p, alternatives outside the eligible
+                # set get effective_probability = 0. When no filter is active,
+                # the two values coincide. Surfacing both lets the panel
+                # report a "model preference margin" (top-1 vs top-2 by
+                # `probability`) and a "sampling margin" (top-1 vs top-2 by
+                # `effective_probability` over the eligible set) — necessary
+                # because under filtering the panel was previously showing a
+                # margin that didn't reflect the actual decision.
+                is_filtered = bool(
+                    (top_k_param is not None and top_k_param > 0)
+                    or (top_p_param is not None and top_p_param < 1.0)
+                )
                 alternatives = []
                 cumulative = 0.0
                 selected_in_top = False
@@ -2184,10 +2199,13 @@ async def analyze_research_attention(request: Dict[str, Any], authenticated: boo
                     cumulative += prob
                     if idx == next_token_id:
                         selected_in_top = True
+                    eff = probs_filtered[idx].item()
                     alternatives.append({
                         "token": token_text,
                         "token_id": idx,
                         "probability": prob,
+                        "effective_probability": eff,
+                        "is_eligible": eff > 0.0,
                         "raw_probability": raw_probs[idx].item(),  # T=1 probability for comparison
                         "log_probability": math.log(prob) if prob > 0 else float('-inf'),
                         "cumulative_probability": cumulative,
@@ -2197,6 +2215,7 @@ async def analyze_research_attention(request: Dict[str, Any], authenticated: boo
                 # If selected token is not in top-N, add it with its actual probability
                 if not selected_in_top:
                     selected_prob = probs[next_token_id].item()
+                    selected_eff = probs_filtered[next_token_id].item()
                     selected_raw_prob = raw_probs[next_token_id].item()
                     selected_log_prob = log_probs[next_token_id].item()
                     selected_logit = raw_logits[next_token_id].item()
@@ -2207,6 +2226,8 @@ async def analyze_research_attention(request: Dict[str, Any], authenticated: boo
                         "token": next_token_text,
                         "token_id": next_token_id,
                         "probability": selected_prob,
+                        "effective_probability": selected_eff,
+                        "is_eligible": selected_eff > 0.0,
                         "raw_probability": selected_raw_prob,  # T=1 probability for comparison
                         "log_probability": selected_log_prob,
                         "cumulative_probability": None,  # Not in sequence
@@ -2223,14 +2244,19 @@ async def analyze_research_attention(request: Dict[str, Any], authenticated: boo
                             "is_selected_outlier": True
                         })
 
-                # Build sampling metadata
+                # Build sampling metadata. `is_filtered` and `eligible_count`
+                # let the frontend decide whether to show one or two
+                # distributions in the panel.
+                eligible_count = int((probs_filtered > 0).sum().item())
                 sampling_metadata = {
                     "temperature": temperature,
                     "top_k": top_k_param,
                     "top_p": top_p_param,
                     "greedy_token_id": greedy_token_id,
                     "greedy_token": greedy_token,
-                    "was_greedy": next_token_id == greedy_token_id
+                    "was_greedy": next_token_id == greedy_token_id,
+                    "is_filtered": is_filtered,
+                    "eligible_count": eligible_count,
                 }
 
                 token_alternatives_by_step.append({
@@ -3098,6 +3124,14 @@ async def analyze_research_attention_stream(request: Dict[str, Any], authenticat
                     log_probs = torch.nn.functional.log_softmax(logits, dim=-1)
                     top_probs = torch.exp(log_probs[top_indices])
 
+                    # See non-SSE site for the rationale behind emitting both
+                    # `probability` (pre-filter model preference) and
+                    # `effective_probability` (post-filter, what the sampler
+                    # drew from). When no filter is active the two coincide.
+                    is_filtered = bool(
+                        (top_k_param is not None and top_k_param > 0)
+                        or (top_p_param is not None and top_p_param < 1.0)
+                    )
                     alternatives = []
                     cumulative = 0.0
                     selected_in_top = False
@@ -3106,10 +3140,13 @@ async def analyze_research_attention_stream(request: Dict[str, Any], authenticat
                         cumulative += prob
                         if idx == next_token_id:
                             selected_in_top = True
+                        eff = probs_filtered[idx].item()
                         alternatives.append({
                             "token": token_text,
                             "token_id": idx,
                             "probability": prob,
+                            "effective_probability": eff,
+                            "is_eligible": eff > 0.0,
                             "raw_probability": raw_probs[idx].item(),  # T=1 probability for comparison
                             "log_probability": math_module.log(prob) if prob > 0 else float('-inf'),
                             "cumulative_probability": cumulative,
@@ -3119,6 +3156,7 @@ async def analyze_research_attention_stream(request: Dict[str, Any], authenticat
                     # If selected token is not in top-N, add it with its actual probability
                     if not selected_in_top:
                         selected_prob = probs[next_token_id].item()
+                        selected_eff = probs_filtered[next_token_id].item()
                         selected_raw_prob = raw_probs[next_token_id].item()
                         selected_log_prob = log_probs[next_token_id].item()
                         selected_logit = raw_logits[next_token_id].item()
@@ -3129,6 +3167,8 @@ async def analyze_research_attention_stream(request: Dict[str, Any], authenticat
                             "token": next_token_text,
                             "token_id": next_token_id,
                             "probability": selected_prob,
+                            "effective_probability": selected_eff,
+                            "is_eligible": selected_eff > 0.0,
                             "raw_probability": selected_raw_prob,  # T=1 probability for comparison
                             "log_probability": selected_log_prob,
                             "cumulative_probability": None,
@@ -3146,13 +3186,16 @@ async def analyze_research_attention_stream(request: Dict[str, Any], authenticat
                             })
 
                     # Build sampling metadata
+                    eligible_count = int((probs_filtered > 0).sum().item())
                     sampling_metadata = {
                         "temperature": temperature,
                         "top_k": top_k_param,
                         "top_p": top_p_param,
                         "greedy_token_id": greedy_token_id,
                         "greedy_token": greedy_token,
-                        "was_greedy": next_token_id == greedy_token_id
+                        "was_greedy": next_token_id == greedy_token_id,
+                        "is_filtered": is_filtered,
+                        "eligible_count": eligible_count,
                     }
 
                     # --- Margin computation and stability classification ---