Update index.html

index.html

@@ -215,13 +215,86 @@
     }catch(e){ console.error(e); showError(e); setStatus('Error'); } finally{ busy=false; }
   }
 
-  // Embedding Viewer
+  // Embedding Viewer (with cache)
   function getEmbState(){ if(!EmbCache[currentModel]) EmbCache[currentModel]={coords:null,neigh:null,ids:null,baseDrawn:false}; return EmbCache[currentModel]; }
-  async function ensureEmbeddings(){ const emb=getEmbState(); if(emb.coords&&emb.neigh&&emb.ids) return emb; const files=MODELS[currentModel].emb; emb.coords=await fetch(files.coords).then(r=>r.json()); emb.neigh=await fetch(files.neigh).then(r=>r.json()); emb.ids=new Set(Object.keys(emb.coords).map(Number)); emb.baseDrawn=false; return emb; }
+  async function ensureEmbeddings(){
+    const emb=getEmbState();
+    if(emb.coords&&emb.neigh&&emb.ids) return emb;
+    const files=MODELS[currentModel].emb;
+    emb.coords=await fetch(files.coords).then(r=>r.json());
+    emb.neigh =await fetch(files.neigh ).then(r=>r.json());
+    emb.ids   =new Set(Object.keys(emb.coords).map(Number));
+    emb.baseDrawn=false;
+    return emb;
+  }
+
   function getBounds(coords){ const pts=Object.values(coords); let minX=Infinity,minY=Infinity,maxX=-Infinity,maxY=-Infinity; for(const [x,y] of pts){ if(x<minX)minX=x; if(y<minY)minY=y; if(x>maxX)maxX=x; if(y>maxY)maxY=y; } return {minX,minY,maxX,maxY}; }
   function makeToXY(coords){ const {minX,minY,maxX,maxY}=getBounds(coords); const pad=18,w=canvas.width-pad*2,h=canvas.height-pad*2; return([x,y])=>{const nx=(x-minX)/(maxX-minX); const ny=(y-minY)/(maxY-minY); return [pad+nx*w,pad+(1-ny)*h];}; }
   function drawBase(emb,toXY){ ctx.clearRect(0,0,canvas.width,canvas.height); ctx.fillStyle='#1a2a3a'; for(const id in emb.coords){ const [x,y]=toXY(emb.coords[id]); ctx.beginPath(); ctx.arc(x,y,2,0,Math.PI*2); ctx.fill(); } emb.baseDrawn=true; }
-  async function drawNeighborhood(tokenId){ const emb=await ensureEmbeddings(); if(!emb.ids.has(tokenId)){ embInfo.innerHTML='<span class="warn">Neighborhood unavailable for this token (not in the top-5k set).</span>'; nbrsEl.innerHTML=''; if(!emb.baseDrawn){const toXY=makeToXY(emb.coords); drawBase(emb,toXY);} return; } const toXY=makeToXY(emb.coords); drawBase(emb,toXY); const targetXY=toXY(emb.coords[tokenId]); const list=emb.neigh.neighbors?.[tokenId]||[]; ctx.fillStyle='#93c5fd'; for(const [nid] of list){ const pt=emb.coords[nid]; if(!pt) continue; const [x,y]=toXY(pt); ctx.beginPath(); ctx.arc(x,y,3.4,0,Math.PI*2); ctx.fill(); } ctx.fillStyle='#22d3ee'; ctx.beginPath(); ctx.arc(targetXY[0],targetXY[1],4.8,0,Math.PI*2); ctx.fill(); nbrsEl.innerHTML=''; embInfo.textContent='Nearest neighbors:'; for(const [nid,sim] of list.slice(0,18)){ const tok=decodeId(nid); const b=document.createElement('div'); b.className='chip'; b.textContent=`${tok}  ${(sim*100).toFixed(1)}%`; nbrsEl.appendChild(b); } }
+
+  // ---- Fallback-aware neighborhood lookup ----
+  async function drawNeighborhood(tokenId){
+    const emb = await ensureEmbeddings();
+
+    // If direct match is missing, try space-prefixed and lowercase variants.
+    let id = tokenId;
+    if (!emb.ids.has(id)) {
+      try {
+        const decoded = decodeId(tokenId);
+        // try " space + token"
+        const sp = await tokenizer(' ' + decoded, { add_special_tokens:false });
+        const cand1 = Array.isArray(sp.input_ids) ? sp.input_ids.at(-1) : null;
+        if (cand1 != null && emb.ids.has(cand1)) id = cand1;
+      } catch {}
+    }
+    if (!emb.ids.has(id)) {
+      try {
+        const low = await tokenizer(decodeId(tokenId).toLowerCase(), { add_special_tokens:false });
+        const cand2 = Array.isArray(low.input_ids) ? low.input_ids.at(-1) : null;
+        if (cand2 != null && emb.ids.has(cand2)) id = cand2;
+      } catch {}
+    }
+
+    // Still not found: show message and (re)draw base if needed
+    if (!emb.ids.has(id)) {
+      embInfo.innerHTML = '<span class="warn">Neighborhood unavailable for this token (not in the current map).</span>';
+      nbrsEl.innerHTML = '';
+      if (!emb.baseDrawn) {
+        const toXY = makeToXY(emb.coords);
+        drawBase(emb, toXY);
+      }
+      return;
+    }
+
+    const toXY = makeToXY(emb.coords);
+    drawBase(emb, toXY); // refresh for clarity
+
+    const targetXY = toXY(emb.coords[id]);
+    const list = (emb.neigh.neighbors?.[id]) || [];
+
+    // neighbors
+    ctx.fillStyle = '#93c5fd';
+    for (const [nid] of list){
+      const pt = emb.coords[nid];
+      if (!pt) continue;
+      const [x,y] = toXY(pt);
+      ctx.beginPath(); ctx.arc(x, y, 3.4, 0, Math.PI*2); ctx.fill();
+    }
+    // target
+    ctx.fillStyle = '#22d3ee';
+    ctx.beginPath(); ctx.arc(targetXY[0], targetXY[1], 4.8, 0, Math.PI*2); ctx.fill();
+
+    // chips
+    nbrsEl.innerHTML = '';
+    embInfo.textContent = 'Nearest neighbors:';
+    for (const [nid, sim] of list.slice(0,18)){
+      const tok = decodeId(nid);
+      const b = document.createElement('div');
+      b.className = 'chip';
+      b.textContent = `${tok}  ${(sim*100).toFixed(1)}%`;
+      nbrsEl.appendChild(b);
+    }
+  }
 
   // Events
   let debounceId;