Add per-token alternatives + hover tooltip UI

server.py

@@ -97,26 +97,81 @@ class WordPredictionRequest(BaseModel):
 @app.post("/api/chat")
 def chat_generate(req: ChatRequest):
     try:
+        # Build prompt and run generation while requesting per-step scores
         prompt = (
             "<|im_start|>system\nYou are a helpful AI assistant.<|im_end|>\n"
             f"<|im_start|>user\n{req.message}<|im_end|>\n"
             "<|im_start|>assistant\n"
         )
         inputs = chat_tokenizer(prompt, return_tensors="pt").to(chat_model.device)
+
+        # Generate deterministically (greedy) while returning scores for each generated step
         outputs = chat_model.generate(
             **inputs,
             max_new_tokens=req.max_new_tokens,
             temperature=req.temperature,
-            do_sample=True,
-            top_p=0.9,
+            do_sample=False,
+            output_scores=True,
+            return_dict_in_generate=True,
             eos_token_id=chat_tokenizer.eos_token_id,
             pad_token_id=chat_tokenizer.eos_token_id,
         )
-        new_tokens = outputs[0][inputs["input_ids"].size(1):]
-        reply = chat_tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
-        if not reply:
-            reply = chat_tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
-        return {"success": True, "response": reply}
+
+        # Full sequence and newly generated token ids
+        sequence = outputs.sequences[0]
+        start_idx = inputs["input_ids"].size(1)
+        generated_ids = sequence[start_idx:].tolist()
+
+        # Decode the full reply
+        reply = chat_tokenizer.decode(generated_ids, skip_special_tokens=True).strip()
+
+        # Prepare per-token alternatives using the per-step logits/scores
+        tokens_info = []
+        # outputs.scores is a tuple with one entry per generated step
+        if hasattr(outputs, "scores") and outputs.scores is not None:
+            for i, logits in enumerate(outputs.scores):
+                # logits shape: (batch_size, vocab_size)
+                probs = torch.softmax(logits[0], dim=-1)
+                chosen_id = generated_ids[i]
+
+                # Get top-k (we ask for 6 and drop the chosen token if present)
+                topk = torch.topk(probs, k=6)
+                alts = []
+                for idx, val in zip(topk.indices.tolist(), topk.values.tolist()):
+                    if idx == chosen_id:
+                        continue
+                    alts.append({
+                        "id": idx,
+                        "token": chat_tokenizer.decode([idx], skip_special_tokens=True).strip(),
+                        "probability": float(val)
+                    })
+                    if len(alts) >= 5:
+                        break
+
+                # Fallback: if not enough alts, sample additional highest-prob tokens
+                if len(alts) < 5:
+                    # get full topk of vocab (expensive but rare for short max_new_tokens)
+                    fallback_topk = torch.topk(probs, k=10)
+                    for idx, val in zip(fallback_topk.indices.tolist(), fallback_topk.values.tolist()):
+                        if idx == chosen_id:
+                            continue
+                        if any(a["id"] == idx for a in alts):
+                            continue
+                        alts.append({
+                            "id": idx,
+                            "token": chat_tokenizer.decode([idx], skip_special_tokens=True).strip(),
+                            "probability": float(val)
+                        })
+                        if len(alts) >= 5:
+                            break
+
+                tokens_info.append({
+                    "id": chosen_id,
+                    "token": chat_tokenizer.decode([chosen_id], skip_special_tokens=True).strip(),
+                    "alternatives": alts
+                })
+
+        return {"success": True, "response": reply, "tokens": tokens_info}
     except Exception as e:
         return {"success": False, "error": str(e)}
 

static/css/styles.css

@@ -409,4 +409,49 @@ textarea:focus {
     background-color: rgba(220, 38, 38, 0.2);
     border-left: 4px solid rgb(220, 38, 38);
     margin: 0 auto;
-}
+}
+
+/* Token hover alternatives */
+.generated-text {
+    display: inline-block;
+    line-height: 1.6;
+}
+.generated-token {
+    display: inline-block;
+    padding: 2px 4px;
+    margin-right: 1px;
+    border-radius: 4px;
+    cursor: pointer;
+    color: var(--text-light);
+}
+.generated-token:hover {
+    background: rgba(139,92,246,0.12);
+}
+.alt-tooltip {
+    position: absolute;
+    display: none;
+    min-width: 160px;
+    background: linear-gradient(180deg, #111827, #0b1220);
+    color: var(--text-light);
+    border: 1px solid rgba(139,92,246,0.18);
+    border-radius: 8px;
+    padding: 8px;
+    box-shadow: 0 8px 24px rgba(2,6,23,0.6);
+    z-index: 2000;
+}
+.alt-title {
+    font-weight: 600;
+    font-size: 0.9rem;
+    margin-bottom: 6px;
+    opacity: 0.9;
+}
+.alt-row {
+    display: flex;
+    justify-content: space-between;
+    gap: 8px;
+    padding: 6px 6px;
+    border-radius: 6px;
+}
+.alt-row:hover { background: rgba(255,255,255,0.02); }
+.alt-token { color: var(--text-light); }
+.alt-prob { color: var(--accent-gray); font-size: 0.85rem; }

static/js/main.js

@@ -148,8 +148,84 @@ async function sendMessage() {
         const data = await response.json();
         const botMessage = document.createElement('div');
         botMessage.className = 'message assistant';
-        botMessage.textContent = data.response || 'Sorry, I could not process your request.';
-        chatOutput.appendChild(botMessage);
+
+        // If the server returned per-token info, render tokens individually so we
+        // can show alternative tokens on hover. Otherwise, fall back to plain text.
+        if (data.tokens && Array.isArray(data.tokens) && data.tokens.length > 0) {
+            const frag = document.createDocumentFragment();
+            const wrapper = document.createElement('div');
+            wrapper.className = 'generated-text';
+
+            data.tokens.forEach((t, idx) => {
+                const span = document.createElement('span');
+                span.className = 'generated-token';
+                span.setAttribute('data-token-index', idx);
+                span.textContent = t.token || '';
+                // store alternatives on the element for quick access
+                span._alternatives = t.alternatives || [];
+                wrapper.appendChild(span);
+            });
+
+            frag.appendChild(wrapper);
+            botMessage.appendChild(frag);
+            chatOutput.appendChild(botMessage);
+
+            // Tooltip element for showing alternatives
+            let tooltip = document.getElementById('alt-tooltip');
+            if (!tooltip) {
+                tooltip = document.createElement('div');
+                tooltip.id = 'alt-tooltip';
+                tooltip.className = 'alt-tooltip';
+                document.body.appendChild(tooltip);
+            }
+
+            // Attach hover listeners
+            wrapper.querySelectorAll('.generated-token').forEach(el => {
+                el.addEventListener('mouseenter', (ev) => {
+                    const alts = el._alternatives || [];
+                    if (!alts.length) return;
+                    // build tooltip html
+                    tooltip.innerHTML = '';
+                    const title = document.createElement('div');
+                    title.className = 'alt-title';
+                    title.textContent = 'Alternatives';
+                    tooltip.appendChild(title);
+
+                    alts.forEach(a => {
+                        const row = document.createElement('div');
+                        row.className = 'alt-row';
+                        const tok = document.createElement('span');
+                        tok.className = 'alt-token';
+                        tok.textContent = a.token || '';
+                        const prob = document.createElement('span');
+                        prob.className = 'alt-prob';
+                        prob.textContent = `${(a.probability * 100).toFixed(2)}%`;
+                        row.appendChild(tok);
+                        row.appendChild(prob);
+                        // click to insert token into input (optional UX)
+                        row.addEventListener('click', () => {
+                            const chatInput = document.getElementById('chat-input');
+                            insertAtCursor(chatInput, a.token || '');
+                        });
+                        tooltip.appendChild(row);
+                    });
+
+                    // Position tooltip near the hovered token
+                    const rect = el.getBoundingClientRect();
+                    tooltip.style.display = 'block';
+                    tooltip.style.left = `${rect.left + window.scrollX}px`;
+                    tooltip.style.top = `${rect.bottom + window.scrollY + 6}px`;
+                });
+
+                el.addEventListener('mouseleave', () => {
+                    const tooltip = document.getElementById('alt-tooltip');
+                    if (tooltip) tooltip.style.display = 'none';
+                });
+            });
+        } else {
+            botMessage.textContent = data.response || 'Sorry, I could not process your request.';
+            chatOutput.appendChild(botMessage);
+        }
 
         // Clear input
         input.value = '';
@@ -317,4 +393,17 @@ document.getElementById('summary-input').addEventListener('keypress', (e) => {
         e.preventDefault();
         generateSummary();
     }
-});
+});
+
+// Helper to insert text at the cursor position for input/textarea
+function insertAtCursor(el, text) {
+    if (!el) return;
+    const start = typeof el.selectionStart === 'number' ? el.selectionStart : el.value.length;
+    const end = typeof el.selectionEnd === 'number' ? el.selectionEnd : start;
+    const before = el.value.substring(0, start);
+    const after = el.value.substring(end);
+    el.value = before + text + after;
+    const pos = before.length + text.length;
+    el.selectionStart = el.selectionEnd = pos;
+    el.focus();
+}