Update app.py

app.py

@@ -4,30 +4,39 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
 # ---- Model config ----
-MODEL_NAME = "gpt2"           # e.g., "distilgpt2", "gpt2", "gpt2-medium"
+MODEL_NAME = "gpt2"  # e.g., "distilgpt2", "gpt2-medium"
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 tok = AutoTokenizer.from_pretrained(MODEL_NAME)
 model = AutoModelForCausalLM.from_pretrained(MODEL_NAME).to(device)
 model.eval()
 
+EPS = 1e-9  # tie tolerance
+
+def safe_exp(x: float) -> str:
+    # Pretty string even for big magnitudes
+    try:
+        return f"{math.exp(x):.6e}"
+    except OverflowError:
+        return "∞ (overflow)"
+    except Exception:
+        return "—"
+
+def is_finite(x: float) -> bool:
+    return x is not None and math.isfinite(x)
+
 def seq_logprob(context: str, candidate: str, assume_leading_space: bool, show_topk: int):
     """
-    Return (total_logprob, step_detail_text, token_list)
-    Computes P(candidate | context) via chain rule over the candidate tokens.
+    Compute log P(candidate | context) via chain rule over tokens.
+    Returns (total_logprob, detail_text, token_list, num_tokens).
     """
-    if not context.strip():
-        return None, "Please enter context.", []
-
-    # Helpful for GPT-2 BPE “word starts”
     cand_text = (" " + candidate) if assume_leading_space else candidate
-
     with torch.no_grad():
         ctx_ids = tok.encode(context, return_tensors="pt").to(device)
         cand_ids = tok.encode(cand_text, add_special_tokens=False)
 
         if len(cand_ids) == 0:
-            return None, "Candidate tokenized to empty sequence (check spacing).", []
+            return None, "Candidate tokenized to empty sequence (check spacing).", [], 0
 
         total_logprob = 0.0
         step_lines = []
@@ -38,8 +47,8 @@ def seq_logprob(context: str, candidate: str, assume_leading_space: bool, show_t
             outputs = model(input_ids=input_ids)
             logits = outputs.logits[:, -1, :]
             logprobs = torch.log_softmax(logits, dim=-1)
-            token_logprob = logprobs[0, t_id].item()
-            total_logprob += token_logprob
+            token_lp = logprobs[0, t_id].item()
+            total_logprob += token_lp
 
             tok_str = tok.decode([t_id])
             token_texts.append(tok_str)
@@ -49,126 +58,147 @@ def seq_logprob(context: str, candidate: str, assume_leading_space: bool, show_t
                 tops = ", ".join([f"{repr(tok.decode([int(idx)]))}:{val.item():.2f}"
                                   for idx, val in zip(topk_idx[0], topk_vals[0])])
                 step_lines.append(
-                    f"Step {i+1}: token={repr(tok_str)}  logprob={token_logprob:.6f}  "
-                    f"prob={math.exp(token_logprob):.6e}\n  top-{show_topk}: {tops}"
+                    f"Step {i+1}: token={repr(tok_str)}  logprob={token_lp:.6f}  "
+                    f"prob={math.exp(token_lp):.6e}\n  top-{show_topk}: {tops}"
                 )
             else:
                 step_lines.append(
-                    f"Step {i+1}: token={repr(tok_str)}  logprob={token_logprob:.6f}  "
-                    f"prob={math.exp(token_logprob):.6e}"
+                    f"Step {i+1}: token={repr(tok_str)}  logprob={token_lp:.6f}  "
+                    f"prob={math.exp(token_lp):.6e}"
                 )
 
             # teacher-forcing: append the true token to continue conditioning
             input_ids = torch.cat([input_ids, torch.tensor([[t_id]], device=device)], dim=1)
 
-        detail_text = "\n".join(step_lines)
-        return total_logprob, detail_text, token_texts
+        return total_logprob, "\n".join(step_lines), token_texts, len(cand_ids)
 
-def compare_candidates(context, cand1, cand2, assume_space, topk):
-    # Basic input checks
-    errs = []
+def compare_candidates(context, candA, candB, assume_space, topk, use_len_norm):
+    # Basic checks
+    errors = []
     if not context.strip():
-        errs.append("Please enter a context.")
-    if not cand1.strip():
-        errs.append("Please enter Candidate A.")
-    if not cand2.strip():
-        errs.append("Please enter Candidate B.")
-    if errs:
-        return (
-            f"<div style='color:#b00020;font-weight:600'>{' '.join(errs)}</div>",
-            "", "", "", "", ""
-        )
-
-    # Compute log-probs
-    logp1, details1, toks1 = seq_logprob(context, cand1, assume_space, topk)
-    logp2, details2, toks2 = seq_logprob(context, cand2, assume_space, topk)
-
-    if logp1 is None or logp2 is None:
-        return (
-            "<div style='color:#b00020;font-weight:600'>Tokenization error. Check inputs.</div>",
-            details1, details2, "", "", ""
-        )
+        errors.append("Please enter a context.")
+    if not candA.strip():
+        errors.append("Please enter Candidate A.")
+    if not candB.strip():
+        errors.append("Please enter Candidate B.")
+    if errors:
+        msg = " ".join(errors)
+        return (f"<div style='color:#b00020;font-weight:600'>{msg}</div>",
+                "", "", "", "", "")
+
+    # Compute
+    lpA, detA, toksA, nA = seq_logprob(context, candA, assume_space, topk)
+    lpB, detB, toksB, nB = seq_logprob(context, candB, assume_space, topk)
+
+    # Validate numbers
+    if not (is_finite(lpA) and is_finite(lpB)):
+        return ("<div style='color:#b00020;font-weight:600'>Numerical issue (NaN/Inf). "
+                "Try shorter context, smaller model (e.g., distilgpt2), or disable length normalization.</div>",
+                "", "", "", "", "")
+
+    # Optionally length-normalize (per-token average log-prob)
+    # Note: odds under length-normalization are "per-token odds", not whole-sequence odds.
+    if use_len_norm:
+        if nA == 0 or nB == 0:
+            return ("<div style='color:#b00020;font-weight:600'>Empty tokenization. "
+                    "Check spacing or turn off 'assume leading space'.</div>",
+                    "", "", "", "", "")
+        scoreA = lpA / nA
+        scoreB = lpB / nB
+        label_suffix = " (per-token)"
+    else:
+        scoreA = lpA
+        scoreB = lpB
+        label_suffix = ""
 
-    # Summaries for each candidate
-    def make_summary(label, cand, logp, toks):
-        seq_prob = math.exp(logp)
-        return (
-            f"**{label}**: {cand}\n\n"
-            f"Tokenization: {toks}\n"
-            f"Total logprob: {logp:.6f}\n"
-            f"Sequence probability: {seq_prob:.6e}"
-        )
+    diff = scoreA - scoreB  # log-odds if unnormalized; log per-token odds otherwise
 
-    summary1 = make_summary("Candidate A", cand1, logp1, toks1)
-    summary2 = make_summary("Candidate B", cand2, logp2, toks2)
+    # Winner logic with proper tie handling
+    if abs(diff) <= EPS:
+        winner = "Tie"
+    elif diff > 0:
+        winner = "Candidate A"
+    else:
+        winner = "Candidate B"
 
-    # Ratio, odds, and winner
-    log_odds = logp1 - logp2                # log(P(A)/P(B))
-    # Cap extreme ratios for display; still show exact log-odds
-    try:
-        ratio = math.exp(log_odds)
-        ratio_str = f"{ratio:.6e}"
-    except OverflowError:
-        ratio_str = "∞ (overflow)"
-    winner = "Candidate A" if logp1 > logp2 else ("Tie" if abs(log_odds) < 1e-12 else "Candidate B")
+    ratio_str = safe_exp(diff)
 
+    # Colors
     if winner == "Candidate A":
         win_color = "#166534"  # green
     elif winner == "Candidate B":
         win_color = "#1d4ed8"  # blue
     else:
-        win_color = "#92400e"  # amber (tie)
+        win_color = "#92400e"  # amber
 
+    # Headline
     headline = (
         f"<div style='padding:14px;border-radius:12px;background:#f8fafc;"
         f"border:1px solid #e2e8f0;margin-bottom:10px'>"
-        f"<div style='font-size:20px;font-weight:800;color:{win_color};'>Winner: {winner}</div>"
+        f"<div style='font-size:20px;font-weight:800;color:{win_color};'>Winner: {winner}{label_suffix}</div>"
         f"<div style='margin-top:6px;font-size:16px;'>"
-        f"Odds (A/B) = <b>{ratio_str}</b> &nbsp;|&nbsp; "
-        f"log-odds = <b>{log_odds:.6f}</b>"
+        f"Odds A/B{label_suffix} = <b>{ratio_str}</b> &nbsp;|&nbsp; "
+        f"log-odds A−B{label_suffix} = <b>{diff:.6f}</b>"
         f"</div>"
         f"<div style='margin-top:6px;color:#475569'>"
-        f"(Odds &gt; 1 means A is more probable; &lt; 1 means B is more probable.)"
+        f"(Odds &gt; 1 ⇒ A more probable; &lt; 1 ⇒ B more probable. "
+        f"{'Per-token uses average log-prob.' if use_len_norm else 'Whole-sequence comparison.'})"
         f"</div></div>"
     )
 
-    return headline, summary1, details1, summary2, details2, ""
+    def summarize(label, cand, lp, toks, n):
+        return (
+            f"**{label}**: {cand}\n\n"
+            f"Tokenization: {toks}\n"
+            f"Total logprob: {lp:.6f}\n"
+            f"Sequence probability: {math.exp(lp):.6e}\n"
+            f"Tokens: {n}"
+        )
+
+    sumA = summarize("Candidate A", candA, lpA, toksA, nA)
+    sumB = summarize("Candidate B", candB, lpB, toksB, nB)
 
-with gr.Blocks(title="Two-Candidate Next-Token Probability Comparator") as demo:
+    return headline, sumA, detA, sumB, detB, ""
+
+def swap(a, b):
+    return b, a
+
+with gr.Blocks(title="Two-Candidate Next-Token Comparator (Robust)") as demo:
     gr.Markdown(
-        "# Two-Candidate Next-Word/Token Probability\n"
-        "Given a **context**, compare the conditional probabilities of **two candidate continuations**.\n"
-        "- Uses a pretrained causal LM (default: GPT-2). No fine-tuning.\n"
-        "- Works at the **token** level; multi-token “words” are handled via the chain rule.\n"
-        "- The **Winner** is the higher-probability candidate; we also show the **odds ratio (A/B)** and log-odds."
+        "# Two-Candidate Next-Word/Token Probability (Robust)\n"
+        "Compare **P(A|context)** vs **P(B|context)** from a pretrained causal LM (no fine-tuning).\n"
+        "- Proper tie handling and numerical guards.\n"
+        "- Optional **length normalization** (per-token).\n"
+        "- Use **Swap** to sanity-check symmetry."
     )
-
     with gr.Row():
-        context = gr.Textbox(label="Context (prompt)", lines=6, placeholder="Paste your prior text here...")
-
+        context = gr.Textbox(label="Context (prompt)", lines=6, placeholder="Paste prior text here...")
     with gr.Row():
-        cand1 = gr.Textbox(label="Candidate A (follow-up)")
-        cand2 = gr.Textbox(label="Candidate B (follow-up)")
-
+        candA = gr.Textbox(label="Candidate A (follow-up)")
+        candB = gr.Textbox(label="Candidate B (follow-up)")
     with gr.Row():
-        assume_space = gr.Checkbox(
-            value=True,
-            label="Assume leading space before candidates (helps align with word starts in GPT-2 tokenization)"
-        )
+        assume_space = gr.Checkbox(value=True, label="Assume leading space before candidates (useful for GPT-2 tokenization)")
         topk = gr.Slider(0, 20, value=5, step=1, label="Show top-k alternatives (per token step)")
+        use_len_norm = gr.Checkbox(value=False, label="Use length normalization (average log-prob per token)")
+    with gr.Row():
+        btn_compare = gr.Button("Compare", variant="primary")
+        btn_swap = gr.Button("Swap A ↔ B")
 
-    btn = gr.Button("Compare")
     winner_html = gr.HTML()
-    summary1 = gr.Markdown()
-    details1 = gr.Textbox(label="Candidate A — step-by-step", lines=10)
-    summary2 = gr.Markdown()
-    details2 = gr.Textbox(label="Candidate B — step-by-step", lines=10)
+    summaryA = gr.Markdown()
+    detailsA = gr.Textbox(label="Candidate A — step-by-step", lines=10)
+    summaryB = gr.Markdown()
+    detailsB = gr.Textbox(label="Candidate B — step-by-step", lines=10)
     _hidden = gr.Textbox(visible=False)
 
-    btn.click(
+    btn_compare.click(
         fn=compare_candidates,
-        inputs=[context, cand1, cand2, assume_space, topk],
-        outputs=[winner_html, summary1, details1, summary2, details2, _hidden]
+        inputs=[context, candA, candB, assume_space, topk, use_len_norm],
+        outputs=[winner_html, summaryA, detailsA, summaryB, detailsB, _hidden]
+    )
+
+    btn_swap.click(
+        fn=swap, inputs=[candA, candB], outputs=[candA, candB]
     )
 
 demo.launch()