Implement generation mode

app.py

@@ -1,26 +1,33 @@
 from enum import Enum
 from pathlib import Path
+from typing import Hashable
 
 import streamlit as st
 import streamlit.components.v1 as components
 import numpy as np
 import torch
 import torch.nn.functional as F
-from transformers import AutoModelForCausalLM, AutoTokenizer, BatchEncoding
+from transformers import AutoModelForCausalLM, AutoTokenizer, BatchEncoding, GPT2LMHeadModel, PreTrainedTokenizer
 
 root_dir = Path(__file__).resolve().parent
 highlighted_text_component = components.declare_component(
     "highlighted_text", path=root_dir / "highlighted_text" / "build"
 )
 
-def get_windows_batched(examples: BatchEncoding, window_len: int, stride: int = 1, pad_id: int = 0) -> BatchEncoding:
+def get_windows_batched(
+    examples: BatchEncoding,
+    window_len: int,
+    start: int = 0,
+    stride: int = 1,
+    pad_id: int = 0
+) -> BatchEncoding:
     return BatchEncoding({
         k: [
             t[i][j : j + window_len] + [
                 pad_id if k in ["input_ids", "labels"] else 0
             ] * (j + window_len - len(t[i]))
             for i in range(len(examples["input_ids"]))
-            for j in range(0, len(examples["input_ids"][i]) - 1, stride)
+            for j in range(start, len(examples["input_ids"][i]), stride)
         ]
         for k, t in examples.items()
     })
@@ -78,13 +85,12 @@ if not compact_layout:
         """
     )
 
-generation_mode = False
-# st.radio("Mode", ["Standard", "Generation"], horizontal=True) == "Generation"
-# st.caption(
-#     "In standard mode, we analyze the model's predictions on the input text. "
-#     "In generation mode, we generate a continuation of the input text "
-#     "and visualize the contributions of different contexts to each generated token."
-# )
+generation_mode = st.radio("Mode", ["Standard", "Generation"], horizontal=True) == "Generation"
+st.caption(
+    "In standard mode, we analyze the model's predictions on the input text. "
+    "In generation mode, we generate a continuation of the input text (prompt) "
+    "and visualize the contributions of different contexts to each generated token."
+)
 
 model_name = st.selectbox("Model", ["distilgpt2", "gpt2", "EleutherAI/gpt-neo-125m"])
 metric_name = st.radio(
@@ -113,6 +119,13 @@ window_len = st.select_slider(
 max_tokens = int(MAX_MEM / (multiplier * window_len) - window_len)
 max_tokens = min(max_tokens, 4096)
 
+max_new_tokens = None
+if generation_mode:
+    max_new_tokens = st.slider(
+        "Max. number of generated tokens",
+        min_value=8, max_value=min(1024, max_tokens), value=min(128, max_tokens)
+    )
+
 DEFAULT_TEXT = """
 We present context length probing, a novel explanation technique for causal
 language models, based on tracking the predictions of a model as a function of the length of
@@ -124,22 +137,21 @@ dependencies.
 """.replace("\n", " ").strip()
 
 text = st.text_area(
-    f"Input text (≤\u2009{max_tokens} tokens)",
+    f"Prompt" if generation_mode else "Input text (≤\u2009{max_tokens} tokens)",
     st.session_state.get("input_text", DEFAULT_TEXT),
     key="input_text",
 )
 
-if tokenizer.eos_token:
-    text += tokenizer.eos_token
 inputs = tokenizer([text])
 [input_ids] = inputs["input_ids"]
-inputs["labels"] = [[*input_ids[1:], tokenizer.eos_token_id]]
-num_user_tokens = len(input_ids) - (1 if tokenizer.eos_token else 0)
+label_ids = [*input_ids[1:], tokenizer.eos_token_id]
+inputs["labels"] = [label_ids]
+num_user_tokens = len(input_ids)
 
 if num_user_tokens < 1:
     st.error("Please enter at least one token.", icon="🚨")
     st.stop()
-if num_user_tokens > max_tokens:
+if not generation_mode and num_user_tokens > max_tokens:
     st.error(
         f"Your input has {num_user_tokens} tokens. Please enter at most {max_tokens} tokens "
         f"or try reducing the window size.",
@@ -150,53 +162,109 @@ if num_user_tokens > max_tokens:
 with st.spinner("Loading model…"):
     model = st.cache_resource(AutoModelForCausalLM.from_pretrained, show_spinner=False)(model_name)
 
-window_len = min(window_len, len(input_ids))
+@torch.inference_mode()
+def get_logprobs(model, inputs, metric):
+    logprobs = []
+    batch_size = 8
+    num_items = len(inputs["input_ids"])
+    pbar = st.progress(0)
+    for i in range(0, num_items, batch_size):
+        pbar.progress(i / num_items, f"{i}/{num_items}")
+        batch = {k: v[i:i + batch_size] for k, v in inputs.items()}
+        batch_logprobs = model(**batch).logits.log_softmax(dim=-1).to(torch.float16)
+        if metric != "KL divergence":
+            batch_logprobs = torch.gather(
+                batch_logprobs, dim=-1, index=batch["labels"][..., None]
+            )
+        logprobs.append(batch_logprobs)
+    logprobs = torch.cat(logprobs, dim=0)
+    pbar.empty()
+    return logprobs
 
 @torch.inference_mode()
-def get_logprobs(_model, _inputs):
-    return _model(**_inputs).logits.log_softmax(dim=-1).to(torch.float16)
+def generate(model, inputs, metric, window_len, max_new_tokens):
+    assert metric == "NLL loss"
+    start = max(0, inputs["input_ids"].shape[1] - window_len + 1)
+    inputs_window = {k: v[:, start:] for k, v in inputs.items()}
+    del inputs_window["labels"]
+
+    new_ids, logprobs = [], []
+    eos_idx = None
+    pbar = st.progress(0)
+    max_steps = max_new_tokens + window_len - 1
+    for i in range(max_steps):
+        pbar.progress(i / max_steps, f"{i}/{max_steps}")
+        inputs_window["attention_mask"] = torch.ones_like(inputs_window["input_ids"], dtype=torch.long)
+        logprobs_window = model(**inputs_window).logits.log_softmax(dim=-1).squeeze(0)
+        if eos_idx is None:
+            next_token = torch.multinomial(logprobs_window[-1].exp(), num_samples=1).item()
+            if next_token == tokenizer.eos_token_id or i >= max_new_tokens - 1:
+                eos_idx = i
+        else:
+            next_token = tokenizer.eos_token_id
+        new_ids.append(next_token)
+
+        inputs_window["input_ids"] = torch.cat([inputs_window["input_ids"], torch.tensor([[next_token]])], dim=1)
+        if inputs_window["input_ids"].shape[1] > window_len:
+            inputs_window["input_ids"] = inputs_window["input_ids"][:, 1:]
+        if logprobs_window.shape[0] == window_len:
+            logprobs.append(
+                logprobs_window[torch.arange(window_len), inputs_window["input_ids"].squeeze(0)]
+            )
+
+        if eos_idx is not None and i - eos_idx >= window_len - 1:
+            break
+    pbar.empty()
+
+    return torch.as_tensor(new_ids[:eos_idx + 1]), torch.stack(logprobs)[:, :, None]
 
-@st.cache_data(show_spinner=False)
 @torch.inference_mode()
-def run_context_length_probing(_model, _tokenizer, _inputs, window_len, metric, cache_key):
+def run_context_length_probing(
+    _model: GPT2LMHeadModel,
+    _tokenizer: PreTrainedTokenizer,
+    _inputs: dict[str, torch.Tensor],
+    window_len: int,
+    metric: str,
+    generation_mode: bool,
+    max_new_tokens: int,
+    cache_key: Hashable
+):
     del cache_key
 
-    inputs_sliding = get_windows_batched(
-        _inputs,
-        window_len=window_len,
-        pad_id=_tokenizer.eos_token_id
-    ).convert_to_tensors("pt")
+    [input_ids] = _inputs["input_ids"]
+    [label_ids] = _inputs["labels"]
 
-    logprobs = []
     with st.spinner("Running model…"):
-        batch_size = 8
-        num_items = len(inputs_sliding["input_ids"])
-        pbar = st.progress(0)
-        for i in range(0, num_items, batch_size):
-            pbar.progress(i / num_items, f"{i}/{num_items}")
-            batch = {k: v[i:i + batch_size] for k, v in inputs_sliding.items()}
-            batch_logprobs = get_logprobs(
-                _model,
-                batch,
-                #cache_key=(model_name, batch["input_ids"].cpu().numpy().tobytes())
+        if generation_mode:
+            new_ids, logprobs = generate(
+                model=_model,
+                inputs=_inputs.convert_to_tensors("pt"), 
+                metric=metric,
+                window_len=window_len,
+                max_new_tokens=max_new_tokens
             )
-            batch_labels = batch["labels"]
-            if metric != "KL divergence":
-                batch_logprobs = torch.gather(
-                    batch_logprobs, dim=-1, index=batch_labels[..., None]
-                )
-            logprobs.append(batch_logprobs)
-        logprobs = torch.cat(logprobs, dim=0)
-        pbar.empty()
+            output_ids = [*input_ids, *new_ids]
+            window_len = logprobs.shape[1]
+        else:
+            window_len = min(window_len, len(input_ids))
+            inputs_sliding = get_windows_batched(
+                _inputs,
+                window_len=window_len,
+                start=0,
+                pad_id=_tokenizer.eos_token_id
+            ).convert_to_tensors("pt")
+            logprobs = get_logprobs(model=model, inputs=inputs_sliding, metric=metric)
+            output_ids = [*input_ids, label_ids[-1]]
+        num_tgt_tokens = logprobs.shape[0]
 
     with st.spinner("Computing scores…"):
         logprobs = logprobs.permute(1, 0, 2)
         logprobs = F.pad(logprobs, (0, 0, 0, window_len, 0, 0), value=torch.nan)
         logprobs = logprobs.view(-1, logprobs.shape[-1])[:-window_len]
-        logprobs = logprobs.view(window_len, len(input_ids) + window_len - 2, logprobs.shape[-1])
+        logprobs = logprobs.view(window_len, num_tgt_tokens + window_len - 1, logprobs.shape[-1])
 
         if metric == "NLL loss":
-            scores = nll_score(logprobs=logprobs, labels=input_ids[1:])
+            scores = nll_score(logprobs=logprobs, labels=label_ids)
         elif metric == "KL divergence":
             scores = kl_div_score(logprobs)
         del logprobs  # possibly destroyed by the score computation to save memory
@@ -206,17 +274,29 @@ def run_context_length_probing(_model, _tokenizer, _inputs, window_len, metric,
         scores /= scores.abs().max(dim=1, keepdim=True).values + 1e-6
         scores = scores.to(torch.float16)
 
-    return scores
+        if generation_mode:
+            scores = F.pad(scores, (0, 0, max(0, len(input_ids) - window_len + 1), 0), value=0.)
+
+    return output_ids, scores
 
-scores = run_context_length_probing(
+if not generation_mode:
+    run_context_length_probing = st.cache_data(run_context_length_probing, show_spinner=False)
+
+output_ids, scores = run_context_length_probing(
     _model=model,
     _tokenizer=tokenizer,
     _inputs=inputs,
     window_len=window_len,
     metric=metric_name,
+    generation_mode=generation_mode,
+    max_new_tokens=max_new_tokens,
     cache_key=(model_name, text),
 )
-tokens = ids_to_readable_tokens(tokenizer, input_ids)
+tokens = ids_to_readable_tokens(tokenizer, output_ids)
 
 st.markdown('<label style="font-size: 14px;">Output</label>', unsafe_allow_html=True)
-highlighted_text_component(tokens=tokens, scores=scores.tolist())
+highlighted_text_component(
+    tokens=tokens,
+    scores=scores.tolist(),
+    prefix_len=len(input_ids) if generation_mode else 0
+)

highlighted_text/build/asset-manifest.json

@@ -1,20 +1,20 @@
 {
   "files": {
-    "main.css": "./static/css/main.b2f89d68.chunk.css",
-    "main.js": "./static/js/main.d475cbb2.chunk.js",
-    "main.js.map": "./static/js/main.d475cbb2.chunk.js.map",
+    "main.css": "./static/css/main.a84f16ea.chunk.css",
+    "main.js": "./static/js/main.5f2b5265.chunk.js",
+    "main.js.map": "./static/js/main.5f2b5265.chunk.js.map",
     "runtime-main.js": "./static/js/runtime-main.e6b0ae4c.js",
     "runtime-main.js.map": "./static/js/runtime-main.e6b0ae4c.js.map",
     "static/js/2.ce130e37.chunk.js": "./static/js/2.ce130e37.chunk.js",
     "static/js/2.ce130e37.chunk.js.map": "./static/js/2.ce130e37.chunk.js.map",
     "index.html": "./index.html",
-    "static/css/main.b2f89d68.chunk.css.map": "./static/css/main.b2f89d68.chunk.css.map",
+    "static/css/main.a84f16ea.chunk.css.map": "./static/css/main.a84f16ea.chunk.css.map",
     "static/js/2.ce130e37.chunk.js.LICENSE.txt": "./static/js/2.ce130e37.chunk.js.LICENSE.txt"
   },
   "entrypoints": [
     "static/js/runtime-main.e6b0ae4c.js",
     "static/js/2.ce130e37.chunk.js",
-    "static/css/main.b2f89d68.chunk.css",
-    "static/js/main.d475cbb2.chunk.js"
+    "static/css/main.a84f16ea.chunk.css",
+    "static/js/main.5f2b5265.chunk.js"
   ]
 }

highlighted_text/build/index.html

@@ -1 +1 @@
-<!doctype html><html lang="en"><head><title>Highlighted text component</title><meta charset="UTF-8"/><meta name="viewport" content="width=device-width,initial-scale=1"/><meta name="description" content="Highlighted text"/><link href="./static/css/main.b2f89d68.chunk.css" rel="stylesheet"></head><body><div id="root"></div><script>!function(e){function t(t){for(var n,l,a=t[0],p=t[1],i=t[2],c=0,s=[];c<a.length;c++)l=a[c],Object.prototype.hasOwnProperty.call(o,l)&&o[l]&&s.push(o[l][0]),o[l]=0;for(n in p)Object.prototype.hasOwnProperty.call(p,n)&&(e[n]=p[n]);for(f&&f(t);s.length;)s.shift()();return u.push.apply(u,i||[]),r()}function r(){for(var e,t=0;t<u.length;t++){for(var r=u[t],n=!0,a=1;a<r.length;a++){var p=r[a];0!==o[p]&&(n=!1)}n&&(u.splice(t--,1),e=l(l.s=r[0]))}return e}var n={},o={1:0},u=[];function l(t){if(n[t])return n[t].exports;var r=n[t]={i:t,l:!1,exports:{}};return e[t].call(r.exports,r,r.exports,l),r.l=!0,r.exports}l.m=e,l.c=n,l.d=function(e,t,r){l.o(e,t)||Object.defineProperty(e,t,{enumerable:!0,get:r})},l.r=function(e){"undefined"!=typeof Symbol&&Symbol.toStringTag&&Object.defineProperty(e,Symbol.toStringTag,{value:"Module"}),Object.defineProperty(e,"__esModule",{value:!0})},l.t=function(e,t){if(1&t&&(e=l(e)),8&t)return e;if(4&t&&"object"==typeof e&&e&&e.__esModule)return e;var r=Object.create(null);if(l.r(r),Object.defineProperty(r,"default",{enumerable:!0,value:e}),2&t&&"string"!=typeof e)for(var n in e)l.d(r,n,function(t){return e[t]}.bind(null,n));return r},l.n=function(e){var t=e&&e.__esModule?function(){return e.default}:function(){return e};return l.d(t,"a",t),t},l.o=function(e,t){return Object.prototype.hasOwnProperty.call(e,t)},l.p="./";var a=this.webpackJsonpstreamlit_component_template=this.webpackJsonpstreamlit_component_template||[],p=a.push.bind(a);a.push=t,a=a.slice();for(var i=0;i<a.length;i++)t(a[i]);var f=p;r()}([])</script><script src="./static/js/2.ce130e37.chunk.js"></script><script src="./static/js/main.d475cbb2.chunk.js"></script></body></html>
+<!doctype html><html lang="en"><head><title>Highlighted text component</title><meta charset="UTF-8"/><meta name="viewport" content="width=device-width,initial-scale=1"/><meta name="description" content="Highlighted text"/><link href="./static/css/main.a84f16ea.chunk.css" rel="stylesheet"></head><body><div id="root"></div><script>!function(e){function t(t){for(var n,l,a=t[0],p=t[1],i=t[2],c=0,s=[];c<a.length;c++)l=a[c],Object.prototype.hasOwnProperty.call(o,l)&&o[l]&&s.push(o[l][0]),o[l]=0;for(n in p)Object.prototype.hasOwnProperty.call(p,n)&&(e[n]=p[n]);for(f&&f(t);s.length;)s.shift()();return u.push.apply(u,i||[]),r()}function r(){for(var e,t=0;t<u.length;t++){for(var r=u[t],n=!0,a=1;a<r.length;a++){var p=r[a];0!==o[p]&&(n=!1)}n&&(u.splice(t--,1),e=l(l.s=r[0]))}return e}var n={},o={1:0},u=[];function l(t){if(n[t])return n[t].exports;var r=n[t]={i:t,l:!1,exports:{}};return e[t].call(r.exports,r,r.exports,l),r.l=!0,r.exports}l.m=e,l.c=n,l.d=function(e,t,r){l.o(e,t)||Object.defineProperty(e,t,{enumerable:!0,get:r})},l.r=function(e){"undefined"!=typeof Symbol&&Symbol.toStringTag&&Object.defineProperty(e,Symbol.toStringTag,{value:"Module"}),Object.defineProperty(e,"__esModule",{value:!0})},l.t=function(e,t){if(1&t&&(e=l(e)),8&t)return e;if(4&t&&"object"==typeof e&&e&&e.__esModule)return e;var r=Object.create(null);if(l.r(r),Object.defineProperty(r,"default",{enumerable:!0,value:e}),2&t&&"string"!=typeof e)for(var n in e)l.d(r,n,function(t){return e[t]}.bind(null,n));return r},l.n=function(e){var t=e&&e.__esModule?function(){return e.default}:function(){return e};return l.d(t,"a",t),t},l.o=function(e,t){return Object.prototype.hasOwnProperty.call(e,t)},l.p="./";var a=this.webpackJsonpstreamlit_component_template=this.webpackJsonpstreamlit_component_template||[],p=a.push.bind(a);a.push=t,a=a.slice();for(var i=0;i<a.length;i++)t(a[i]);var f=p;r()}([])</script><script src="./static/js/2.ce130e37.chunk.js"></script><script src="./static/js/main.5f2b5265.chunk.js"></script></body></html>

highlighted_text/build/static/css/{main.b2f89d68.chunk.css → main.a84f16ea.chunk.css}

@@ -1,2 +1,2 @@
-body{padding:0;margin:0;font-family:"Source Sans Pro",sans-serif;font-size:1rem;line-height:1.4}.highlighted-text,.status-bar{color:#000;background-color:#fff}.highlighted-text{border:1px solid #d2d2d2;border-radius:0 0 5px 5px;padding:4px;cursor:pointer}.highlighted-text .token.active{outline:1px solid #444}.highlighted-text.frozen .token{opacity:.75}.highlighted-text.frozen .token.active,.highlighted-text.frozen .token.context{opacity:1}.highlighted-text.frozen .token.context{-webkit-text-decoration:#999 underline;text-decoration:#999 underline}.status-bar{min-height:1.4em;border:1px solid #d2d2d2;border-bottom:none;border-radius:5px 5px 0 0;padding:3px 4px}.status-bar .token{border:1px solid #aaa;border-radius:2px;background:#f5f5f5;padding:1px;margin:1px 2px 1px 1px}
-/*# sourceMappingURL=main.b2f89d68.chunk.css.map */
+body{padding:0;margin:0;font-family:"Source Sans Pro",sans-serif;font-size:1rem;line-height:1.4}.highlighted-text,.status-bar{color:#000;background-color:#fff}.highlighted-text{border:1px solid #d2d2d2;border-radius:0 0 5px 5px;padding:4px;cursor:pointer}.highlighted-text .token.prefix~.token:not(.prefix){color:#2563eb}.highlighted-text .token.active{outline:1px solid #444}.status-bar{min-height:1.4em;border:1px solid #d2d2d2;border-bottom:none;border-radius:5px 5px 0 0;padding:3px 4px}.status-bar .token{border:1px solid #aaa;border-radius:2px;background:#f5f5f5;padding:1px;margin:1px 2px 1px 1px}
+/*# sourceMappingURL=main.a84f16ea.chunk.css.map */

highlighted_text/build/static/css/main.a84f16ea.chunk.css.map

@@ -0,0 +1 @@
+{"version":3,"sources":["webpack://src/index.scss"],"names":[],"mappings":"AAAA,KACI,SAAA,CACA,QAAA,CACA,wCAAA,CACA,cAAA,CACA,eAAA,CAGJ,8BACI,UAAA,CACA,qBAAA,CAGJ,kBACI,wBAAA,CACA,yBAAA,CACA,WAAA,CACA,cAAA,CAEA,oDACI,aAAA,CAGJ,gCACI,sBAAA,CAIR,YACI,gBAAA,CAEA,wBAAA,CAAA,kBAAA,CACA,yBAAA,CACA,eAAA,CAEA,mBACI,qBAAA,CACA,iBAAA,CACA,kBAAA,CACA,WAAA,CAEA,sBAAA","file":"main.a84f16ea.chunk.css","sourcesContent":["body {\n    padding: 0;\n    margin: 0;\n    font-family: \"Source Sans Pro\", sans-serif;\n    font-size: 1rem;\n    line-height: 1.4;\n}\n\n.highlighted-text, .status-bar {\n    color: black;\n    background-color: white;\n}\n\n.highlighted-text {\n    border: 1px solid #d2d2d2;\n    border-radius: 0 0 5px 5px;\n    padding: 4px;\n    cursor: pointer;\n\n    .token.prefix ~ .token:not(.prefix) {\n        color: #2563eb;\n    }\n\n    .token.active {\n        outline: 1px solid #444;\n    }\n}\n\n.status-bar {\n    min-height: 1.4em;\n    border: 1px solid #d2d2d2;\n    border-bottom: none;\n    border-radius: 5px 5px 0 0;\n    padding: 3px 4px;\n\n    .token {\n        border: 1px solid #aaa;\n        border-radius: 2px;\n        background: #f5f5f5;\n        padding: 1px;\n        margin: 1px;\n        margin-right: 2px;\n    }\n}\n"]}

highlighted_text/build/static/css/main.b2f89d68.chunk.css.map

@@ -1 +0,0 @@
-{"version":3,"sources":["webpack://src/index.scss"],"names":[],"mappings":"AAAA,KACI,SAAA,CACA,QAAA,CACA,wCAAA,CACA,cAAA,CACA,eAAA,CAGJ,8BACI,UAAA,CACA,qBAAA,CAGJ,kBACI,wBAAA,CACA,yBAAA,CACA,WAAA,CACA,cAAA,CAEA,gCACI,sBAAA,CAIA,gCACI,WAAA,CAEA,+EACI,SAAA,CAGJ,wCACI,sCAAA,CAAA,8BAAA,CAMhB,YACI,gBAAA,CAEA,wBAAA,CAAA,kBAAA,CACA,yBAAA,CACA,eAAA,CAEA,mBACI,qBAAA,CACA,iBAAA,CACA,kBAAA,CACA,WAAA,CAEA,sBAAA","file":"main.b2f89d68.chunk.css","sourcesContent":["body {\n    padding: 0;\n    margin: 0;\n    font-family: \"Source Sans Pro\", sans-serif;\n    font-size: 1rem;\n    line-height: 1.4;\n}\n\n.highlighted-text, .status-bar {\n    color: black;\n    background-color: white;\n}\n\n.highlighted-text {\n    border: 1px solid #d2d2d2;\n    border-radius: 0 0 5px 5px;\n    padding: 4px;\n    cursor: pointer;\n\n    .token.active {\n        outline: 1px solid #444;\n    }\n\n    &.frozen {\n        .token {\n            opacity: 0.75;\n\n            &.context, &.active {\n                opacity: 1;\n            }\n\n            &.context {\n                text-decoration: #999 underline;\n            }\n        }\n    }\n}\n\n.status-bar {\n    min-height: 1.4em;\n    border: 1px solid #d2d2d2;\n    border-bottom: none;\n    border-radius: 5px 5px 0 0;\n    padding: 3px 4px;\n\n    .token {\n        border: 1px solid #aaa;\n        border-radius: 2px;\n        background: #f5f5f5;\n        padding: 1px;\n        margin: 1px;\n        margin-right: 2px;\n    }\n}\n"]}

highlighted_text/build/static/js/main.5f2b5265.chunk.js

@@ -0,0 +1,2 @@
+(this.webpackJsonpstreamlit_component_template=this.webpackJsonpstreamlit_component_template||[]).push([[0],{27:function(t,e,n){},28:function(t,e,n){"use strict";n.r(e);var a=n(7),s=n.n(a),r=n(18),c=n.n(r),i=n(4),o=n(0),l=n(1),h=n(2),d=n(3),j=n(16),u=n(6),x=function(t){Object(h.a)(n,t);var e=Object(d.a)(n);function n(){var t;Object(o.a)(this,n);for(var a=arguments.length,s=new Array(a),r=0;r<a;r++)s[r]=arguments[r];return(t=e.call.apply(e,[this].concat(s))).state={activeIndex:null,hoverIndex:null,isFrozen:!1},t}return Object(l.a)(n,[{key:"render",value:function(){var t=this,e=this.props.args.tokens,n=this.getScores(),a=this.props.args.prefix_len,s="highlighted-text";this.state.isFrozen&&(s+=" frozen");var r=function(){t.setState({isFrozen:!1})};return Object(u.jsxs)(u.Fragment,{children:[Object(u.jsxs)("div",{className:"status-bar",children:[Object(u.jsxs)("span",{className:this.state.isFrozen?"":" d-none",children:[Object(u.jsx)("i",{className:"fa fa-lock"})," "]},"lock-icon"),null!=this.state.activeIndex?Object(u.jsxs)(u.Fragment,{children:[Object(u.jsx)("strong",{children:"index:"},"index-label")," ",Object(u.jsxs)("span",{children:[this.state.activeIndex," "]},"index")]}):Object(u.jsx)(u.Fragment,{})]},"status-bar"),Object(u.jsx)("div",{className:s,onClick:r,children:e.map((function(e,s){var c="token";t.state&&t.state.activeIndex==s&&(c+=" active"),s<a&&(c+=" prefix");var i={backgroundColor:n[s]>0?"rgba(32, 255, 32, ".concat(n[s],")"):"rgba(255, 32, 32, ".concat(-n[s],")")};return Object(u.jsx)("span",{className:c,style:i,onMouseOver:function(){t.state.isFrozen||t.setState({activeIndex:s}),t.setState({hoverIndex:s})},onClick:r,children:e},s)}))},"text")]})}},{key:"getScores",value:function(){var t=this.props.args.tokens;if(!this.state||null==this.state.activeIndex||this.state.activeIndex<1)return t.map((function(){return 0}));var e=this.props.args.scores,n=this.state.activeIndex-1,a=Math.min(Math.max(0,n),e[n].length),s=e[n].slice(0,a);s.reverse();var r=[].concat(Object(i.a)(Array(Math.max(0,n-s.length)).fill(0)),Object(i.a)(s.map((function(t){return void 0==t||isNaN(t)?0:t}))));return r=[].concat(Object(i.a)(r),Object(i.a)(Array(t.length-r.length).fill(0)))}}]),n}(j.a),b=Object(j.b)(x);n(27);c.a.render(Object(u.jsx)(s.a.StrictMode,{children:Object(u.jsx)(b,{})}),document.getElementById("root"))}},[[28,1,2]]]);
+//# sourceMappingURL=main.5f2b5265.chunk.js.map

highlighted_text/build/static/js/main.5f2b5265.chunk.js.map

@@ -0,0 +1 @@
+{"version":3,"sources":["HighlightedText.tsx","index.tsx"],"names":["HighlightedText","_StreamlitComponentBa","_inherits","_super","_createSuper","_this","_classCallCheck","_len","arguments","length","args","Array","_key","call","apply","concat","state","activeIndex","hoverIndex","isFrozen","_createClass","key","value","_this2","tokens","this","props","scores","getScores","prefixLength","className","onClick","setState","_jsxs","_Fragment","children","_jsx","map","t","i","style","backgroundColor","onMouseOver","allScores","hi","Math","min","max","row","slice","reverse","result","_toConsumableArray","fill","x","undefined","isNaN","StreamlitComponentBase","withStreamlitConnection","ReactDOM","render","React","StrictMode","document","getElementById"],"mappings":"gQAeMA,EAAe,SAAAC,GAAAC,YAAAF,EAAAC,GAAA,IAAAE,EAAAC,YAAAJ,GAAA,SAAAA,IAAA,IAAAK,EAAAC,YAAA,KAAAN,GAAA,QAAAO,EAAAC,UAAAC,OAAAC,EAAA,IAAAC,MAAAJ,GAAAK,EAAA,EAAAA,EAAAL,EAAAK,IAAAF,EAAAE,GAAAJ,UAAAI,GACqD,OADrDP,EAAAF,EAAAU,KAAAC,MAAAX,EAAA,OAAAY,OAAAL,KACVM,MAAQ,CAACC,YAAa,KAAMC,WAAY,KAAMC,UAAU,GAAMd,CAAC,CA6ErE,OA7EoEe,YAAApB,EAAA,EAAAqB,IAAA,SAAAC,MAErE,WAAU,IAADC,EAAA,KACCC,EAAmBC,KAAKC,MAAMhB,KAAa,OAC3CiB,EAAmBF,KAAKG,YACxBC,EAAeJ,KAAKC,MAAMhB,KAAiB,WAE7CoB,EAAY,mBACZL,KAAKT,MAAMG,WACXW,GAAa,WAGjB,IAAMC,EAAU,WACZR,EAAKS,SAAS,CAAEb,UAAU,GAC9B,EAEA,OAAOc,eAAAC,WAAA,CAAAC,SAAA,CACHF,eAAA,OAAKH,UAAU,aAAYK,SAAA,CACvBF,eAAA,QAAMH,UAAWL,KAAKT,MAAMG,SAAW,GAAK,UAAUgB,SAAA,CAAiBC,cAAA,KAAGN,UAAU,eAAiB,MAA1C,aAE7B,MAA1BL,KAAKT,MAAMC,YACXgB,eAAAC,WAAA,CAAAC,SAAA,CACIC,cAAA,UAAAD,SAA0B,UAAd,eAA6B,IAACF,eAAA,QAAAE,SAAA,CAAmBV,KAAKT,MAAMC,YAAY,MAAhC,YAEtDmB,cAAAF,WAAA,MAPsB,cAUhCE,cAAA,OAAKN,UAAWA,EAAWC,QAASA,EAAQI,SAEpCX,EAAOa,KAAI,SAACC,EAAWC,GACnB,IAAIT,EAAY,QACZP,EAAKP,OACDO,EAAKP,MAAMC,aAAesB,IAC1BT,GAAa,WAGjBS,EAAIV,IACJC,GAAa,WAEjB,IAAMU,EAAQ,CACVC,gBACId,EAAOY,GAAK,EAAC,qBAAAxB,OACcY,EAAOY,GAAE,0BAAAxB,QACRY,EAAOY,GAAE,MAS7C,OAAOH,cAAA,QAAcN,UAAWA,EAAWU,MAAOA,EAC9CE,YAPgB,WACXnB,EAAKP,MAAMG,UACZI,EAAKS,SAAS,CAAEf,YAAasB,IAEjChB,EAAKS,SAAS,CAAEd,WAAYqB,GAChC,EAE8BR,QAASA,EAAQI,SAAEG,GAD/BC,EAEtB,KA3ByC,UA+BzD,GAAC,CAAAlB,IAAA,YAAAC,MAED,WACI,IAAME,EAASC,KAAKC,MAAMhB,KAAa,OACvC,IAAKe,KAAKT,OAAmC,MAA1BS,KAAKT,MAAMC,aAAuBQ,KAAKT,MAAMC,YAAc,EAC1E,OAAOO,EAAOa,KAAI,kBAAM,CAAC,IAE7B,IAAMM,EAAwBlB,KAAKC,MAAMhB,KAAa,OAEhD6B,EAAId,KAAKT,MAAMC,YAAc,EAC7B2B,EAAKC,KAAKC,IAAID,KAAKE,IAAI,EAAGR,GAAII,EAAUJ,GAAG9B,QAC3CuC,EAAML,EAAUJ,GAAGU,MAAM,EAAGL,GAClCI,EAAIE,UACJ,IAAIC,EAAM,GAAApC,OAAAqC,YACHzC,MAAMkC,KAAKE,IAAI,EAAGR,EAAIS,EAAIvC,SAAS4C,KAAK,IAAED,YAC1CJ,EAAIX,KAAI,SAACiB,GAAC,YAAUC,GAALD,GAAkBE,MAAMF,GAAK,EAAIA,CAAC,MAGxD,OADAH,EAAM,GAAApC,OAAAqC,YAAOD,GAAMC,YAAKzC,MAAMa,EAAOf,OAAS0C,EAAO1C,QAAQ4C,KAAK,IAEtE,KAACrD,CAAA,CA9EgB,CAASyD,KAiFfC,cAAwB1D,G,MC3FvC2D,IAASC,OACPxB,cAACyB,IAAMC,WAAU,CAAA3B,SACfC,cAACpC,EAAe,MAElB+D,SAASC,eAAe,Q","file":"static/js/main.5f2b5265.chunk.js","sourcesContent":["import {\n    StreamlitComponentBase,\n    withStreamlitConnection,\n} from \"streamlit-component-lib\";\n\ntype HighlightedTextState = {\n    activeIndex: number | null,\n    hoverIndex: number | null,\n    isFrozen: boolean\n};\n\n/**\n * This is a React-based component template. The `render()` function is called\n * automatically when your component should be re-rendered.\n */\nclass HighlightedText extends StreamlitComponentBase<HighlightedTextState> {\n    public state = {activeIndex: null, hoverIndex: null, isFrozen: false};\n\n    render() {\n        const tokens: string[] = this.props.args[\"tokens\"];\n        const scores: number[] = this.getScores();\n        const prefixLength = this.props.args[\"prefix_len\"];\n\n        let className = \"highlighted-text\";\n        if (this.state.isFrozen) {\n            className += \" frozen\";\n        }\n\n        const onClick = () => {\n            this.setState({ isFrozen: false });\n        };\n\n        return <>\n            <div className=\"status-bar\" key=\"status-bar\">\n                <span className={this.state.isFrozen ? \"\" : \" d-none\"} key=\"lock-icon\"><i className=\"fa fa-lock\"></i> </span>\n                {\n                    this.state.activeIndex != null ?\n                    <>\n                        <strong key=\"index-label\">index:</strong> <span key=\"index\">{this.state.activeIndex} </span>\n                    </>\n                    : <></>\n                }\n            </div>\n            <div className={className} onClick={onClick} key=\"text\">\n                {\n                    tokens.map((t: string, i: number) => {\n                        let className = \"token\";\n                        if (this.state) {\n                            if (this.state.activeIndex == i) {\n                                className += \" active\";\n                            }\n                        }\n                        if (i < prefixLength) {\n                            className += \" prefix\";\n                        }\n                        const style = {\n                            backgroundColor:\n                                scores[i] > 0\n                                    ? `rgba(32, 255, 32, ${scores[i]})`\n                                    : `rgba(255, 32, 32, ${-scores[i]})`\n                        };\n\n                        const onMouseOver = () => {\n                            if (!this.state.isFrozen) {\n                                this.setState({ activeIndex: i });\n                            }\n                            this.setState({ hoverIndex: i });\n                        };\n                        return <span key={i} className={className} style={style}\n                            onMouseOver={onMouseOver} onClick={onClick}>{t}</span>;\n                    })\n                }\n            </div>\n        </>;\n    }\n\n    private getScores() {\n        const tokens = this.props.args[\"tokens\"];\n        if (!this.state || this.state.activeIndex == null || this.state.activeIndex < 1) {\n            return tokens.map(() => 0);\n        }\n        const allScores: number[][] = this.props.args[\"scores\"];\n\n        const i = this.state.activeIndex - 1;\n        const hi = Math.min(Math.max(0, i), allScores[i].length);\n        const row = allScores[i].slice(0, hi);\n        row.reverse();\n        let result = [\n            ...Array(Math.max(0, i - row.length)).fill(0),\n            ...row.map((x) => x == undefined || isNaN(x) ? 0 : x)\n        ];\n        result = [...result, ...Array(tokens.length - result.length).fill(0)];\n        return result;\n    }\n}\n\nexport default withStreamlitConnection(HighlightedText);\n","import React from \"react\";\nimport ReactDOM from \"react-dom\";\nimport HighlightedText from \"./HighlightedText\";\nimport \"./index.scss\";\n\nReactDOM.render(\n  <React.StrictMode>\n    <HighlightedText />\n  </React.StrictMode>,\n  document.getElementById(\"root\")\n)\n"],"sourceRoot":""}

highlighted_text/build/static/js/main.d475cbb2.chunk.js

@@ -1,2 +0,0 @@
-(this.webpackJsonpstreamlit_component_template=this.webpackJsonpstreamlit_component_template||[]).push([[0],{27:function(t,e,n){},28:function(t,e,n){"use strict";n.r(e);var a=n(7),s=n.n(a),c=n(18),r=n.n(c),i=n(4),o=n(0),l=n(1),h=n(2),d=n(3),j=n(16),u=n(6),b=function(t){Object(h.a)(n,t);var e=Object(d.a)(n);function n(){var t;Object(o.a)(this,n);for(var a=arguments.length,s=new Array(a),c=0;c<a;c++)s[c]=arguments[c];return(t=e.call.apply(e,[this].concat(s))).state={activeIndex:null,hoverIndex:null,isFrozen:!1},t}return Object(l.a)(n,[{key:"render",value:function(){var t=this,e=this.props.args.tokens,n=this.getScores(),a="highlighted-text";this.state.isFrozen&&(a+=" frozen");var s=function(){t.setState({isFrozen:!1})};return Object(u.jsxs)(u.Fragment,{children:[Object(u.jsxs)("div",{className:"status-bar",children:[Object(u.jsxs)("span",{className:this.state.isFrozen?"":" d-none",children:[Object(u.jsx)("i",{className:"fa fa-lock"})," "]},"lock-icon"),null!=this.state.activeIndex?Object(u.jsxs)(u.Fragment,{children:[Object(u.jsx)("strong",{children:"index:"},"index-label")," ",Object(u.jsxs)("span",{children:[this.state.activeIndex," "]},"index")]}):Object(u.jsx)(u.Fragment,{})]},"status-bar"),Object(u.jsx)("div",{className:a,onClick:s,children:e.map((function(e,a){var c="token";t.state&&t.state.activeIndex==a&&(c+=" active");var r={backgroundColor:n[a]>0?"rgba(32, 255, 32, ".concat(n[a],")"):"rgba(255, 32, 32, ".concat(-n[a],")")};return Object(u.jsx)("span",{className:c,style:r,onMouseOver:function(){t.state.isFrozen||t.setState({activeIndex:a}),t.setState({hoverIndex:a})},onClick:s,children:e},a)}))},"text")]})}},{key:"getScores",value:function(){var t=this.props.args.tokens;if(!this.state||null==this.state.activeIndex||this.state.activeIndex<1)return t.map((function(){return 0}));var e=this.props.args.scores,n=this.state.activeIndex-1,a=Math.min(Math.max(0,n),e[n].length),s=e[n].slice(0,a);s.reverse();var c=[].concat(Object(i.a)(Array(Math.max(0,n-s.length)).fill(0)),Object(i.a)(s.map((function(t){return void 0==t||isNaN(t)?0:t}))));return c=[].concat(Object(i.a)(c),Object(i.a)(Array(t.length-c.length).fill(0)))}}]),n}(j.a),v=Object(j.b)(b);n(27);r.a.render(Object(u.jsx)(s.a.StrictMode,{children:Object(u.jsx)(v,{})}),document.getElementById("root"))}},[[28,1,2]]]);
-//# sourceMappingURL=main.d475cbb2.chunk.js.map

highlighted_text/build/static/js/main.d475cbb2.chunk.js.map

@@ -1 +0,0 @@
-{"version":3,"sources":["HighlightedText.tsx","index.tsx"],"names":["HighlightedText","_StreamlitComponentBa","_inherits","_super","_createSuper","_this","_classCallCheck","_len","arguments","length","args","Array","_key","call","apply","concat","state","activeIndex","hoverIndex","isFrozen","_createClass","key","value","_this2","tokens","this","props","scores","getScores","className","onClick","setState","_jsxs","_Fragment","children","_jsx","map","t","i","style","backgroundColor","onMouseOver","allScores","hi","Math","min","max","row","slice","reverse","result","_toConsumableArray","fill","x","undefined","isNaN","StreamlitComponentBase","withStreamlitConnection","ReactDOM","render","React","StrictMode","document","getElementById"],"mappings":"gQAeMA,EAAe,SAAAC,GAAAC,YAAAF,EAAAC,GAAA,IAAAE,EAAAC,YAAAJ,GAAA,SAAAA,IAAA,IAAAK,EAAAC,YAAA,KAAAN,GAAA,QAAAO,EAAAC,UAAAC,OAAAC,EAAA,IAAAC,MAAAJ,GAAAK,EAAA,EAAAA,EAAAL,EAAAK,IAAAF,EAAAE,GAAAJ,UAAAI,GACqD,OADrDP,EAAAF,EAAAU,KAAAC,MAAAX,EAAA,OAAAY,OAAAL,KACVM,MAAQ,CAACC,YAAa,KAAMC,WAAY,KAAMC,UAAU,GAAMd,CAAC,CAyErE,OAzEoEe,YAAApB,EAAA,EAAAqB,IAAA,SAAAC,MAErE,WAAU,IAADC,EAAA,KACCC,EAAmBC,KAAKC,MAAMhB,KAAa,OAC3CiB,EAAmBF,KAAKG,YAE1BC,EAAY,mBACZJ,KAAKT,MAAMG,WACXU,GAAa,WAGjB,IAAMC,EAAU,WACZP,EAAKQ,SAAS,CAAEZ,UAAU,GAC9B,EAEA,OAAOa,eAAAC,WAAA,CAAAC,SAAA,CACHF,eAAA,OAAKH,UAAU,aAAYK,SAAA,CACvBF,eAAA,QAAMH,UAAWJ,KAAKT,MAAMG,SAAW,GAAK,UAAUe,SAAA,CAAiBC,cAAA,KAAGN,UAAU,eAAiB,MAA1C,aAE7B,MAA1BJ,KAAKT,MAAMC,YACXe,eAAAC,WAAA,CAAAC,SAAA,CACIC,cAAA,UAAAD,SAA0B,UAAd,eAA6B,IAACF,eAAA,QAAAE,SAAA,CAAmBT,KAAKT,MAAMC,YAAY,MAAhC,YAEtDkB,cAAAF,WAAA,MAPsB,cAUhCE,cAAA,OAAKN,UAAWA,EAAWC,QAASA,EAAQI,SAEpCV,EAAOY,KAAI,SAACC,EAAWC,GACnB,IAAIT,EAAY,QACZN,EAAKP,OACDO,EAAKP,MAAMC,aAAeqB,IAC1BT,GAAa,WAGrB,IAAMU,EAAQ,CACVC,gBACIb,EAAOW,GAAK,EAAC,qBAAAvB,OACcY,EAAOW,GAAE,0BAAAvB,QACRY,EAAOW,GAAE,MAS7C,OAAOH,cAAA,QAAcN,UAAWA,EAAWU,MAAOA,EAC9CE,YAPgB,WACXlB,EAAKP,MAAMG,UACZI,EAAKQ,SAAS,CAAEd,YAAaqB,IAEjCf,EAAKQ,SAAS,CAAEb,WAAYoB,GAChC,EAE8BR,QAASA,EAAQI,SAAEG,GAD/BC,EAEtB,KAxByC,UA4BzD,GAAC,CAAAjB,IAAA,YAAAC,MAED,WACI,IAAME,EAASC,KAAKC,MAAMhB,KAAa,OACvC,IAAKe,KAAKT,OAAmC,MAA1BS,KAAKT,MAAMC,aAAuBQ,KAAKT,MAAMC,YAAc,EAC1E,OAAOO,EAAOY,KAAI,kBAAM,CAAC,IAE7B,IAAMM,EAAwBjB,KAAKC,MAAMhB,KAAa,OAEhD4B,EAAIb,KAAKT,MAAMC,YAAc,EAC7B0B,EAAKC,KAAKC,IAAID,KAAKE,IAAI,EAAGR,GAAII,EAAUJ,GAAG7B,QAC3CsC,EAAML,EAAUJ,GAAGU,MAAM,EAAGL,GAClCI,EAAIE,UACJ,IAAIC,EAAM,GAAAnC,OAAAoC,YACHxC,MAAMiC,KAAKE,IAAI,EAAGR,EAAIS,EAAItC,SAAS2C,KAAK,IAAED,YAC1CJ,EAAIX,KAAI,SAACiB,GAAC,YAAUC,GAALD,GAAkBE,MAAMF,GAAK,EAAIA,CAAC,MAGxD,OADAH,EAAM,GAAAnC,OAAAoC,YAAOD,GAAMC,YAAKxC,MAAMa,EAAOf,OAASyC,EAAOzC,QAAQ2C,KAAK,IAEtE,KAACpD,CAAA,CA1EgB,CAASwD,KA6EfC,cAAwBzD,G,MCvFvC0D,IAASC,OACPxB,cAACyB,IAAMC,WAAU,CAAA3B,SACfC,cAACnC,EAAe,MAElB8D,SAASC,eAAe,Q","file":"static/js/main.d475cbb2.chunk.js","sourcesContent":["import {\n    StreamlitComponentBase,\n    withStreamlitConnection,\n} from \"streamlit-component-lib\";\n\ntype HighlightedTextState = {\n    activeIndex: number | null,\n    hoverIndex: number | null,\n    isFrozen: boolean\n};\n\n/**\n * This is a React-based component template. The `render()` function is called\n * automatically when your component should be re-rendered.\n */\nclass HighlightedText extends StreamlitComponentBase<HighlightedTextState> {\n    public state = {activeIndex: null, hoverIndex: null, isFrozen: false};\n\n    render() {\n        const tokens: string[] = this.props.args[\"tokens\"];\n        const scores: number[] = this.getScores();\n\n        let className = \"highlighted-text\";\n        if (this.state.isFrozen) {\n            className += \" frozen\";\n        }\n\n        const onClick = () => {\n            this.setState({ isFrozen: false });\n        };\n\n        return <>\n            <div className=\"status-bar\" key=\"status-bar\">\n                <span className={this.state.isFrozen ? \"\" : \" d-none\"} key=\"lock-icon\"><i className=\"fa fa-lock\"></i> </span>\n                {\n                    this.state.activeIndex != null ?\n                    <>\n                        <strong key=\"index-label\">index:</strong> <span key=\"index\">{this.state.activeIndex} </span>\n                    </>\n                    : <></>\n                }\n            </div>\n            <div className={className} onClick={onClick} key=\"text\">\n                {\n                    tokens.map((t: string, i: number) => {\n                        let className = \"token\";\n                        if (this.state) {\n                            if (this.state.activeIndex == i) {\n                                className += \" active\";\n                            }\n                        }\n                        const style = {\n                            backgroundColor:\n                                scores[i] > 0\n                                    ? `rgba(32, 255, 32, ${scores[i]})`\n                                    : `rgba(255, 32, 32, ${-scores[i]})`\n                        };\n\n                        const onMouseOver = () => {\n                            if (!this.state.isFrozen) {\n                                this.setState({ activeIndex: i });\n                            }\n                            this.setState({ hoverIndex: i });\n                        };\n                        return <span key={i} className={className} style={style}\n                            onMouseOver={onMouseOver} onClick={onClick}>{t}</span>;\n                    })\n                }\n            </div>\n        </>;\n    }\n\n    private getScores() {\n        const tokens = this.props.args[\"tokens\"];\n        if (!this.state || this.state.activeIndex == null || this.state.activeIndex < 1) {\n            return tokens.map(() => 0);\n        }\n        const allScores: number[][] = this.props.args[\"scores\"];\n\n        const i = this.state.activeIndex - 1;\n        const hi = Math.min(Math.max(0, i), allScores[i].length);\n        const row = allScores[i].slice(0, hi);\n        row.reverse();\n        let result = [\n            ...Array(Math.max(0, i - row.length)).fill(0), \n            ...row.map((x) => x == undefined || isNaN(x) ? 0 : x)\n        ];\n        result = [...result, ...Array(tokens.length - result.length).fill(0)];\n        return result;\n    }\n}\n\nexport default withStreamlitConnection(HighlightedText);\n","import React from \"react\";\nimport ReactDOM from \"react-dom\";\nimport HighlightedText from \"./HighlightedText\";\nimport \"./index.scss\";\n\nReactDOM.render(\n  <React.StrictMode>\n    <HighlightedText />\n  </React.StrictMode>,\n  document.getElementById(\"root\")\n)\n"],"sourceRoot":""}

highlighted_text/src/HighlightedText.tsx

@@ -19,6 +19,7 @@ class HighlightedText extends StreamlitComponentBase<HighlightedTextState> {
     render() {
         const tokens: string[] = this.props.args["tokens"];
         const scores: number[] = this.getScores();
+        const prefixLength: number = this.props.args["prefix_len"];
 
         let className = "highlighted-text";
         if (this.state.isFrozen) {
@@ -49,6 +50,9 @@ class HighlightedText extends StreamlitComponentBase<HighlightedTextState> {
                                 className += " active";
                             }
                         }
+                        if (i < prefixLength) {
+                            className += " prefix";
+                        }
                         const style = {
                             backgroundColor:
                                 scores[i] > 0

highlighted_text/src/index.scss

@@ -17,22 +17,12 @@ body {
     padding: 4px;
     cursor: pointer;
 
-    .token.active {
-        outline: 1px solid #444;
+    .token.prefix ~ .token:not(.prefix) {
+        color: #2563eb;
     }
 
-    &.frozen {
-        .token {
-            opacity: 0.75;
-
-            &.context, &.active {
-                opacity: 1;
-            }
-
-            &.context {
-                text-decoration: #999 underline;
-            }
-        }
+    .token.active {
+        outline: 1px solid #444;
     }
 }