static/script.js

document.getElementById("analyze-button").addEventListener("click", async () => {
    const text = document.getElementById("input-text").value;
    const model = document.getElementById("model-select").value;

    const response = await fetch("/analyze", {
        method: "POST",
        headers: {
            "Content-Type": "application/json"
        },
        body: JSON.stringify({ text, model })
    });

    const data = await response.json();

    const coloredTextDiv = document.getElementById("colored-text");
    coloredTextDiv.innerHTML = "";

    // Always add the first token
    const firstToken = data.tokens[0];
    const firstTokenSpan = document.createElement("span");
    firstTokenSpan.classList.add("token");
    
    // Handle special tokens and regular tokens differently
    if (firstToken === "<s>" || firstToken === "<|endoftext|>") {
        firstTokenSpan.style.backgroundColor = "#808080"; // Gray for special tokens
        firstTokenSpan.textContent = "■";
        tippy(firstTokenSpan, {
            content: "<div><strong>Beginning of Sequence</strong></div>",
            allowHTML: true,
            theme: 'custom',
            placement: 'top',
            interactive: true
        });
    } else {
        // Handle regular first token
        firstTokenSpan.style.backgroundColor = "#808080"; // or any other color you prefer
        firstTokenSpan.textContent = firstToken;
        tippy(firstTokenSpan, {
            content: `<div><strong>First Token</strong></div>`,
            allowHTML: true,
            theme: 'custom',
            placement: 'top',
            interactive: true
        });
    }
    
    coloredTextDiv.appendChild(firstTokenSpan);

    for (let index = 0; index < data.log_probs.length; index++) {
        const token = data.tokens[index + 1];
        const percentile = data.percentiles[index];
        const logProb = data.log_probs[index];
        const topKPredictions = data.top_k_predictions[index];
        const color = getColor(data.log_probs, logProb);

        const tokenSpan = document.createElement("span");
        tokenSpan.classList.add("token");
        tokenSpan.style.backgroundColor = color;

        let displayToken = token;
        let specialTokenDescription = "";

        // Enhanced special token handling
        if (token === "<s>" || token === "<|endoftext|>") {
            displayToken = "■";
            specialTokenDescription = "Beginning of Sequence";
        } else if (token === "</s>" || token === "<|endoftext|>") {
            displayToken = "■";
            specialTokenDescription = "End of Sequence";
        } else if (token === "<0x0A>") {
            displayToken = "■";
            specialTokenDescription = "Newline";
        } else if (token.startsWith("<") && token.endsWith(">")) {
            displayToken = "■";
            specialTokenDescription = "Special Token: " + token;
        } else {
            // Clean up GPT-2 style tokens (Ġ and Ċ)
            displayToken = displayToken
                .replace(/\u2581/g, " ")  // Replace underscore token
                .replace(/Ġ/g, " ")       // Replace GPT-2 space token
                .replace(/Ċ/g, "\n");     // Replace GPT-2 newline token
        }

        tokenSpan.textContent = displayToken;

        let tooltipContent = "";
        if (specialTokenDescription) {
            tooltipContent += `<div style="font-weight: bold; margin-bottom: 8px;">${specialTokenDescription}</div>`;
        }

        tooltipContent += `<div style="font-weight: bold; margin-bottom: 4px;">Top 5 Predictions:</div>`;
        topKPredictions.forEach(pred => {
            let predToken = pred.token;
            if (predToken === "<0x0A>") {
                predToken = "\\n";
            } else if (predToken.startsWith("<") && predToken.endsWith(">")) {
                predToken = "[SPECIAL]";
            } else {
                predToken = predToken
                    .replace(/\u2581/g, " ")
                    .replace(/Ġ/g, " ")
                    .replace(/Ċ/g, "\n");
            }
            tooltipContent += `<div style="padding-left: 8px;">${predToken}: ${pred.log_prob.toFixed(4)}</div>`;
        });
        
        tooltipContent += `<div style="margin-top: 8px; border-top: 1px solid #555; padding-top: 8px;">
            <div><strong>Stats:</strong></div>
            <div style="padding-left: 8px;">Percentile: ${percentile.toFixed(2)}</div>
            <div style="padding-left: 8px;">Log-Likelihood: ${logProb.toFixed(4)}</div>
        </div>`;

        tippy(tokenSpan, {
            content: tooltipContent,
            allowHTML: true,
            theme: 'custom',
            placement: 'top',
            interactive: true
        });

        coloredTextDiv.appendChild(tokenSpan);
        if (token === "<0x0A>") {
          coloredTextDiv.appendChild(document.createElement("br"));
        }
    }

    document.getElementById("joint-log-likelihood").textContent = data.joint_log_likelihood.toFixed(4);
    document.getElementById("average-log-likelihood").textContent = data.average_log_likelihood.toFixed(4);
});

function getColor(allLogProbs, currentLogProb) {
    const minLogProb = Math.min(...allLogProbs);
    const maxLogProb = Math.max(...allLogProbs);

    // Normalize to 0-1 range
    let normalizedLogProb = (currentLogProb - minLogProb) / (maxLogProb - minLogProb);
    normalizedLogProb = Math.max(0, Math.min(1, normalizedLogProb)); // Clamp

    // Optional: Apply a power transformation (adjust the exponent as needed)
    const power = 0.7; // Example: Less than 1 emphasizes differences at lower end
    normalizedLogProb = Math.pow(normalizedLogProb, power);

    const hue = normalizedLogProb * 120; // 0 (red) to 120 (green)
    return `hsl(${hue}, 100%, 50%)`;
}