Update app.py

app.py

@@ -3,98 +3,172 @@ import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
-MODEL_NAME = "gpt2"          # swap to "gpt2-medium" etc. if you like
+# ---- Model config ----
+MODEL_NAME = "gpt2"           # e.g., "distilgpt2", "gpt2", "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()
 
-def next_seq_prob(context, candidate, assume_leading_space, show_topk):
+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.
+    """
     if not context.strip():
-        return "Please enter context.", "", ""
-    if not candidate.strip():
-        return "Please enter a candidate next word/token.", "", ""
+        return None, "Please enter context.", []
 
-    # Optionally prepend a leading space (helps align with GPT-2 BPE “word” starts)
+    # Helpful for GPT-2 BPE “word starts”
     cand_text = (" " + candidate) if assume_leading_space else candidate
 
     with torch.no_grad():
-        # Encode context
         ctx_ids = tok.encode(context, return_tensors="pt").to(device)
-        # Tokenize candidate (no special tokens)
         cand_ids = tok.encode(cand_text, add_special_tokens=False)
+
         if len(cand_ids) == 0:
-            return "Candidate tokenized to empty sequence (check spacing).", "", ""
+            return None, "Candidate tokenized to empty sequence (check spacing).", []
 
         total_logprob = 0.0
-        step_details = []
-
-        # Start from context, then feed each candidate token step-by-step (teacher forcing)
+        step_lines = []
         input_ids = ctx_ids
+        token_texts = []
+
         for i, t_id in enumerate(cand_ids):
             outputs = model(input_ids=input_ids)
-            logits = outputs.logits[:, -1, :]                 # distribution over next token
+            logits = outputs.logits[:, -1, :]
             logprobs = torch.log_softmax(logits, dim=-1)
             token_logprob = logprobs[0, t_id].item()
             total_logprob += token_logprob
 
-            # top-k display
-            topk_vals, topk_idx = torch.topk(logprobs, k=min(show_topk, logprobs.shape[-1]), dim=-1)
-            topk_pairs = [
-                (tok.decode(int(idx)), float(val))
-                for idx, val in zip(topk_idx[0].tolist(), topk_vals[0].tolist())
-            ]
-
-            step_details.append({
-                "step": i+1,
-                "predicted_for": tok.decode([t_id]),
-                "logprob": token_logprob,
-                "prob": math.exp(token_logprob),
-                "topk": topk_pairs
-            })
-
-            # append the true token to continue conditioning
+            tok_str = tok.decode([t_id])
+            token_texts.append(tok_str)
+
+            if show_topk > 0:
+                topk_vals, topk_idx = torch.topk(logprobs, k=min(show_topk, logprobs.shape[-1]), dim=-1)
+                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}"
+                )
+            else:
+                step_lines.append(
+                    f"Step {i+1}: token={repr(tok_str)}  logprob={token_logprob:.6f}  "
+                    f"prob={math.exp(token_logprob):.6e}"
+                )
+
+            # teacher-forcing: append the true token to continue conditioning
             input_ids = torch.cat([input_ids, torch.tensor([[t_id]], device=device)], dim=1)
 
-        seq_prob = math.exp(total_logprob)
-        # Human-friendly note about words vs tokens
-        tokenized_candidate = [tok.decode([i]) for i in cand_ids]
-        summary = (
-            f"Candidate tokenization: {tokenized_candidate}\n"
-            f"Total logprob (chain rule): {total_logprob:.6f}\n"
+        detail_text = "\n".join(step_lines)
+        return total_logprob, detail_text, token_texts
+
+def compare_candidates(context, cand1, cand2, assume_space, topk):
+    # Basic input checks
+    errs = []
+    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, "", "", ""
+        )
+
+    # 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}"
         )
 
-        # Pretty print step details
-        lines = []
-        for d in step_details:
-            lines.append(
-                f"Step {d['step']}: token={repr(d['predicted_for'])}  "
-                f"logprob={d['logprob']:.6f}  prob={d['prob']:.6e}"
-            )
-            if show_topk > 0:
-                tops = ", ".join([f"{repr(tok)}:{lp:.2f}" for tok, lp in d["topk"]])
-                lines.append(f"  top-{show_topk} logprobs: {tops}")
-        detail_text = "\n".join(lines)
-
-        return summary, detail_text, ""
-        
-with gr.Blocks(title="Next-Token Probability (no fine-tuning)") as demo:
-    gr.Markdown("# Next-Token Probability\n"
-                "Compute the probability of a chosen next word/token sequence given a prior text segment.")
+    summary1 = make_summary("Candidate A", cand1, logp1, toks1)
+    summary2 = make_summary("Candidate B", cand2, logp2, toks2)
+
+    # 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")
+
+    if winner == "Candidate A":
+        win_color = "#166534"  # green
+    elif winner == "Candidate B":
+        win_color = "#1d4ed8"  # blue
+    else:
+        win_color = "#92400e"  # amber (tie)
+
+    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='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"</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"</div></div>"
+    )
+
+    return headline, summary1, details1, summary2, details2, ""
+
+with gr.Blocks(title="Two-Candidate Next-Token Probability Comparator") 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."
+    )
+
     with gr.Row():
-        context = gr.Textbox(label="Context (prompt)", lines=6)
+        context = gr.Textbox(label="Context (prompt)", lines=6, placeholder="Paste your prior text here...")
+
     with gr.Row():
-        candidate = gr.Textbox(label="Candidate next word / token sequence")
+        cand1 = gr.Textbox(label="Candidate A (follow-up)")
+        cand2 = gr.Textbox(label="Candidate B (follow-up)")
+
     with gr.Row():
-        assume_space = gr.Checkbox(value=True, label="Assume leading space before candidate (useful for word starts in GPT-2 tokenization)")
-        topk = gr.Slider(0, 20, value=10, step=1, label="Show top-k alternatives (per step)")
-    btn = gr.Button("Compute probability")
-    summary = gr.Textbox(label="Summary", lines=4)
-    details = gr.Textbox(label="Step-by-step (per token)", lines=12)
-    _hidden = gr.Textbox(visible=False)  # placeholder
-
-    btn.click(fn=next_seq_prob, inputs=[context, candidate, assume_space, topk], outputs=[summary, details, _hidden])
+        assume_space = gr.Checkbox(
+            value=True,
+            label="Assume leading space before candidates (helps align with word starts in GPT-2 tokenization)"
+        )
+        topk = gr.Slider(0, 20, value=5, step=1, label="Show top-k alternatives (per token step)")
+
+    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)
+    _hidden = gr.Textbox(visible=False)
+
+    btn.click(
+        fn=compare_candidates,
+        inputs=[context, cand1, cand2, assume_space, topk],
+        outputs=[winner_html, summary1, details1, summary2, details2, _hidden]
+    )
 
 demo.launch()