Update app.py

app.py

@@ -34,17 +34,10 @@ def analyze_text(text, top_k):
 .tt{display:none;position:absolute;bottom:calc(100% + 8px);left:50%;transform:translateX(-50%);
 background:#1a1a2e;color:#eee;padding:14px;border-radius:12px;font-size:13px;z-index:9999;
 box-shadow:0 10px 30px rgba(0,0,0,.35);min-width:220px;max-height:350px;overflow-y:auto}
-
-/* Bridge the gap between token and tooltip with a pseudo-element */
 .tt::after{content:'';position:absolute;top:100%;left:0;width:100%;height:12px}
-
-/* Show on hover */
 .tw:hover .tt{display:block}
-
-/* Show when pinned via click */
 .tw.pinned .tt{display:block}
 .tw.pinned .tk{transform:translateY(-2px);box-shadow:0 4px 14px rgba(0,0,0,.18);border-color:#999;outline:2px solid #7fdbca}
-
 .th{font-weight:700;font-size:14px;color:#7fdbca;border-bottom:1px solid #333;padding-bottom:6px;margin-bottom:6px}
 .tp{color:#ffd700;margin-bottom:8px}
 .at{color:#ff79c6;font-size:10px;text-transform:uppercase;letter-spacing:1px;margin-bottom:4px}
@@ -56,15 +49,11 @@ box-shadow:0 10px 30px rgba(0,0,0,.35);min-width:220px;max-height:350px;overflow
 document.addEventListener('click', function(e) {
     const tk = e.target.closest('.tk');
     const tw = tk ? tk.closest('.tw') : null;
-    // If clicked a token, toggle its pin
     if (tw) {
         const wasPinned = tw.classList.contains('pinned');
-        // Unpin all others
         document.querySelectorAll('.tw.pinned').forEach(el => el.classList.remove('pinned'));
-        // Toggle this one
         if (!wasPinned) tw.classList.add('pinned');
     } else if (!e.target.closest('.tt')) {
-        // Clicked outside any token or tooltip — unpin all
         document.querySelectorAll('.tw.pinned').forEach(el => el.classList.remove('pinned'));
     }
 });
@@ -105,13 +94,48 @@ document.addEventListener('click', function(e) {
     return ''.join(parts)
 
 
+def predict_next(text, num_tokens, temperature):
+    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)
+
+    if max_tokens <= 0:
+        return "Input too long to generate more."
+
+    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()
+
+
 with gr.Blocks(theme=gr.themes.Soft(), css="footer{display:none!important}.main{max-width:960px;margin:auto}") as demo:
-    gr.Markdown("# 🔍 Token Probability Explorer\nPaste text, **hover** to preview or **click** a token to pin its tooltip open. Click elsewhere to dismiss.")
+    gr.Markdown("# 🔍 Token Probability Explorer & Predictor\nPaste text, **hover** to preview or **click** a token to pin its tooltip open. Click elsewhere to dismiss.")
+
+    text_input = gr.Textbox(label="Input Text", placeholder="Paste your text here…", lines=5)
+
     with gr.Row():
-        text_input = gr.Textbox(label="Input Text", placeholder="Paste your text here…", lines=5, scale=4)
-        top_k_input = gr.Number(label="# Alternatives", value=10, minimum=1, maximum=200, step=1, scale=1)
-    btn = gr.Button("Analyze", variant="primary")
-    output = gr.HTML()
-    btn.click(fn=analyze_text, inputs=[text_input, top_k_input], outputs=output)
+        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)
+
+    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)
+
+    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)
 
 demo.launch()