import math
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
# ---- Model config ----
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):
"""
Compute log P(candidate | context) via chain rule over tokens.
Returns (total_logprob, detail_text, token_list, num_tokens).
"""
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).", [], 0
total_logprob = 0.0
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, :]
logprobs = torch.log_softmax(logits, dim=-1)
token_lp = logprobs[0, t_id].item()
total_logprob += token_lp
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_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_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)
return total_logprob, "\n".join(step_lines), token_texts, len(cand_ids)
def compare_candidates(context, candA, candB, assume_space, topk, use_len_norm):
# Basic checks
errors = []
if not context.strip():
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"{msg}
",
"", "", "", "", "")
# 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 ("Numerical issue (NaN/Inf). "
"Try shorter context, smaller model (e.g., distilgpt2), or disable length normalization.
",
"", "", "", "", "")
# 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 ("Empty tokenization. "
"Check spacing or turn off 'assume leading space'.
",
"", "", "", "", "")
scoreA = lpA / nA
scoreB = lpB / nB
label_suffix = " (per-token)"
else:
scoreA = lpA
scoreB = lpB
label_suffix = ""
diff = scoreA - scoreB # log-odds if unnormalized; log per-token odds otherwise
# Winner logic with proper tie handling
if abs(diff) <= EPS:
winner = "Tie"
elif diff > 0:
winner = "Candidate A"
else:
winner = "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
# Headline
headline = (
f""
f"
Winner: {winner}{label_suffix}
"
f"
"
f"Odds A/B{label_suffix} = {ratio_str} | "
f"log-odds A−B{label_suffix} = {diff:.6f}"
f"
"
f"
"
f"(Odds > 1 ⇒ A more probable; < 1 ⇒ B more probable. "
f"{'Per-token uses average log-prob.' if use_len_norm else 'Whole-sequence comparison.'})"
f"