Update app.py

app.py

@@ -94,28 +94,55 @@ document.addEventListener('click', function(e) {
     return ''.join(parts)
 
 
-def predict_next(text, num_tokens, temperature):
+def predict_next(text, num_candidates):
+    num_candidates = max(1, int(num_candidates))
     if not text.strip():
-        return ""
-    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512).to(device)
-    input_len = inputs['input_ids'].shape[1]
-    max_tokens = min(int(num_tokens), 512 - input_len)
+        return "<p style='color:#999;text-align:center;padding:40px'>Enter text and click Predict Next.</p>"
 
-    if max_tokens <= 0:
-        return "Input too long to generate more."
+    tokens = tokenizer.encode(text)
+    if len(tokens) > 512:
+        tokens = tokens[:512]
 
     with torch.no_grad():
-        output_ids = model.generate(
-            **inputs,
-            max_new_tokens=max_tokens,
-            do_sample=True if temperature > 0 else False,
-            temperature=temperature if temperature > 0 else 1.0,
-            top_k=50,
-            pad_token_id=tokenizer.eos_token_id,
-            repetition_penalty=1.1
-        )
-    result = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-    return result[len(text):].strip()
+        input_ids = torch.tensor([tokens]).to(device)
+        logits = model(input_ids).logits[0, -1].cpu()
+
+    probs = torch.softmax(logits, dim=-1)
+    log_probs = torch.log(probs)
+    top_p, top_idx = probs.topk(num_candidates)
+    top_lp = log_probs[top_idx]
+
+    rows = ""
+    for j in range(num_candidates):
+        tok_text = html_lib.escape(tokenizer.decode([top_idx[j].item()]))
+        p = top_p[j].item()
+        lp = top_lp[j].item()
+        bar_width = max(1, int(p * 100))
+        hue = min(p * 120, 120)
+        rows += f"""<tr>
+<td style="padding:6px 12px;font-weight:600;color:#e0e0e0;white-space:nowrap">{j+1}</td>
+<td style="padding:6px 12px;font-family:monospace;font-size:15px;color:#7fdbca;white-space:nowrap">{tok_text}</td>
+<td style="padding:6px 12px;width:100%">
+  <div style="background:hsla({hue},80%,50%,0.25);border-radius:4px;height:22px;width:{bar_width}%;min-width:2px;display:flex;align-items:center;padding-left:6px">
+    <span style="font-size:11px;color:hsl({hue},80%,70%);font-weight:600">{p:.4f}</span>
+  </div>
+</td>
+<td style="padding:6px 12px;font-family:monospace;font-size:13px;color:#888;white-space:nowrap">{lp:.4f}</td>
+</tr>"""
+
+    html = f"""<div style="font-family:'Segoe UI',sans-serif;background:#1a1a2e;border-radius:12px;padding:16px;overflow-x:auto">
+<div style="color:#ff79c6;font-size:11px;text-transform:uppercase;letter-spacing:1px;margin-bottom:10px">
+Top {num_candidates} predicted next tokens</div>
+<table style="width:100%;border-collapse:collapse">
+<thead><tr style="border-bottom:1px solid #333">
+<th style="padding:6px 12px;text-align:left;color:#666;font-size:11px">#</th>
+<th style="padding:6px 12px;text-align:left;color:#666;font-size:11px">TOKEN</th>
+<th style="padding:6px 12px;text-align:left;color:#666;font-size:11px">PROBABILITY</th>
+<th style="padding:6px 12px;text-align:left;color:#666;font-size:11px">LOG PROB</th>
+</tr></thead>
+<tbody>{rows}</tbody>
+</table></div>"""
+    return html
 
 
 with gr.Blocks() as demo:
@@ -125,18 +152,17 @@ with gr.Blocks() as demo:
 
     with gr.Row():
         top_k_input = gr.Number(label="# Alternatives (Analysis)", value=10, minimum=1, maximum=200, step=1)
-        num_tokens_input = gr.Number(label="# Tokens to Predict", value=10, minimum=1, maximum=100, step=1)
-        temperature_input = gr.Slider(label="Temperature (Prediction)", minimum=0.0, maximum=2.0, value=0.7, step=0.05)
+        num_candidates_input = gr.Number(label="# Next Token Candidates", value=10, minimum=1, maximum=200, step=1)
 
     with gr.Row():
         btn_analyze = gr.Button("Analyze", variant="primary")
         btn_predict = gr.Button("Predict Next", variant="secondary")
 
     output_analysis = gr.HTML(label="Analysis Output")
-    output_prediction = gr.Textbox(label="Predicted Continuation", lines=3, interactive=False)
+    output_prediction = gr.HTML(label="Predicted Next Tokens")
 
     btn_analyze.click(fn=analyze_text, inputs=[text_input, top_k_input], outputs=output_analysis)
-    btn_predict.click(fn=predict_next, inputs=[text_input, num_tokens_input, temperature_input], outputs=output_prediction)
+    btn_predict.click(fn=predict_next, inputs=[text_input, num_candidates_input], outputs=output_prediction)
 
 demo.launch(
     server_name="0.0.0.0",