Fix score computation, improve UI, add more models

app.py

@@ -1,3 +1,4 @@
+from enum import Enum
 from pathlib import Path
 
 import streamlit as st
@@ -7,13 +8,10 @@ import torch.nn.functional as F
 from transformers import AutoModelForCausalLM, AutoTokenizer, BatchEncoding
 
 root_dir = Path(__file__).resolve().parent
-_highlighted_text_fn = components.declare_component(
+highlighted_text_component = components.declare_component(
     "highlighted_text", path=root_dir / "highlighted_text" / "build"
 )
 
-def highlighted_text(tokens, scores, key=None):
-    return _highlighted_text_fn(tokens=tokens, scores=scores, key=key, default=0)
-
 def get_windows_batched(examples: BatchEncoding, window_len: int, stride: int = 1, pad_id: int = 0) -> BatchEncoding:
     return BatchEncoding({
         k: [
@@ -43,13 +41,22 @@ def ids_to_readable_tokens(tokenizer, ids, strip_whitespace=False):
             result.append("")
     return result
 
-model_name = st.selectbox("Model", ["distilgpt2"])
+st.header("Context length probing")
+
+with st.form("form"):
+    model_name = st.selectbox("Model", ["distilgpt2", "gpt2"])
+    metric_name = st.selectbox("Metric", ["Cross entropy"])
+
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name)
 
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
+    window_len = st.select_slider("Window size", options=[8, 16, 32, 64, 128, 256, 512, 1024], value=512)
+    text = st.text_area(
+        "Input text",
+        "The complex houses married and single soldiers and their families.",
+    )
 
-window_len = st.select_slider("Window size", options=[8, 16, 32, 64, 128, 256, 512, 1024], value=512)
-text = st.text_area("Input text", "The complex houses married and single soldiers and their families.")
+    st.form_submit_button("Submit")
 
 inputs = tokenizer([text])
 [input_ids] = inputs["input_ids"]
@@ -63,17 +70,16 @@ inputs_sliding = get_windows_batched(
 )
 with torch.inference_mode():
     logits = model(**inputs_sliding.convert_to_tensors("pt")).logits.to(torch.float16)
-    logits = F.pad(logits, (0, 0, 0, window_len - 1, 0, 0), value=torch.nan)
-    logits = logits.view(-1, logits.shape[-1])[:(window_len - 1) * (len(input_ids) + window_len - 2)]
-    logits = logits.view(window_len - 1, len(input_ids) + window_len - 2, logits.shape[-1])
+    logits = logits.permute(1, 0, 2)
+    logits = F.pad(logits, (0, 0, 0, window_len, 0, 0), value=torch.nan)
+    logits = logits.view(-1, logits.shape[-1])[:-window_len]
+    logits = logits.view(window_len, len(input_ids) + window_len - 2, logits.shape[-1])
 
-scores = logits.to(torch.float32).softmax(dim=-1)
+scores = logits.to(torch.float32).log_softmax(dim=-1)
 scores = scores[:, torch.arange(len(input_ids[1:])), input_ids[1:]]
 scores = scores.diff(dim=0).transpose(0, 1)
 scores = scores.nan_to_num()
 scores /= scores.abs().max(dim=1, keepdim=True).values + 1e-9
 scores = scores.to(torch.float16)
-print(scores)
-st.markdown("---")
-highlighted_text(tokens=tokens, scores=scores.tolist())
 
+highlighted_text_component(tokens=tokens, scores=scores.tolist())

highlighted_text/build/asset-manifest.json

@@ -1,7 +1,7 @@
 {
   "files": {
-    "main.js": "./static/js/main.243f5634.chunk.js",
-    "main.js.map": "./static/js/main.243f5634.chunk.js.map",
+    "main.js": "./static/js/main.0c3ca5a6.chunk.js",
+    "main.js.map": "./static/js/main.0c3ca5a6.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",
@@ -12,6 +12,6 @@
   "entrypoints": [
     "static/js/runtime-main.e6b0ae4c.js",
     "static/js/2.ce130e37.chunk.js",
-    "static/js/main.243f5634.chunk.js"
+    "static/js/main.0c3ca5a6.chunk.js"
   ]
 }

highlighted_text/build/index.css

@@ -6,14 +6,11 @@ body {
     line-height: 1.4;
 }
 
-.status-bar {
-    display: none;
-}
-
 .highlighted-text {
     border: 1px solid #d2d2d2;
-    border-radius: 5px 5px 5px 5px;
-    padding: 2px;
+    border-radius: 0 0 5px 5px;
+    background-color: white;
+    padding: 4px;
     cursor: pointer;
 
     .token.active {
@@ -33,4 +30,21 @@ body {
             }
         }
     }
-}
+}
+
+.status-bar {
+    min-height: 1.4em;
+    border: 1px solid #d2d2d2;
+    border-bottom: none;
+    border-radius: 5px 5px 0 0;
+    padding: 3px 4px;
+
+    .token {
+        border: 1px solid #aaa;
+        border-radius: 2px;
+        background: #f5f5f5;
+        padding: 1px;
+        margin: 1px;
+        margin-right: 2px;
+    }
+}

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="theme-color" content="#fff"/><meta name="description" content="Highlighted text"/><link rel="stylesheet" href="index.css"/></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.243f5634.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 rel="stylesheet" href="index.css"/></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.0c3ca5a6.chunk.js"></script></body></html>

highlighted_text/build/static/js/main.0c3ca5a6.chunk.js

@@ -0,0 +1,2 @@
+(this.webpackJsonpstreamlit_component_template=this.webpackJsonpstreamlit_component_template||[]).push([[0],{27:function(t,e,a){"use strict";a.r(e);var n=a(7),s=a.n(n),r=a(18),c=a.n(r),i=a(4),o=a(0),l=a(1),h=a(2),d=a(3),j=a(16),u=a(6),b=function(t){Object(h.a)(a,t);var e=Object(d.a)(a);function a(){var t;Object(o.a)(this,a);for(var n=arguments.length,s=new Array(n),r=0;r<n;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)(a,[{key:"render",value:function(){var t=this,e=this.props.args.tokens,a=this.getScores(),n="highlighted-text";this.state.isFrozen&&(n+=" 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:n,onClick:s,children:e.map((function(e,n){var r="token";t.state&&t.state.activeIndex==n&&(r+=" active");var c={backgroundColor:a[n]>0?"rgba(32, 255, 32, ".concat(a[n],")"):"rgba(255, 32, 32, ".concat(-a[n],")")};return Object(u.jsx)("span",{className:r,style:c,onMouseOver:function(){t.state.isFrozen||t.setState({activeIndex:n}),t.setState({hoverIndex:n})},onClick:s,children:e},n)}))},"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,a=this.state.activeIndex-1,n=Math.min(Math.max(0,a),e[a].length),s=e[a].slice(0,n);s.reverse();var r=[].concat(Object(i.a)(Array(Math.max(0,a-1-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)))}}]),a}(j.a),v=Object(j.b)(b);c.a.render(Object(u.jsx)(s.a.StrictMode,{children:Object(u.jsx)(v,{})}),document.getElementById("root"))}},[[27,1,2]]]);
+//# sourceMappingURL=main.0c3ca5a6.chunk.js.map

highlighted_text/build/static/js/main.0c3ca5a6.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","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":"2OAeMA,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,CAuErE,OAvEoEe,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,aAChC,MAA1BJ,KAAKT,MAAMC,YACZe,eAAAC,WAAA,CAAAC,SAAA,CACIC,cAAA,UAAAD,SAA0B,UAAd,eAA6B,IAACF,eAAA,QAAAE,SAAA,CAAmBT,KAAKT,MAAMC,YAAY,MAAhC,YAEtDkB,cAAAF,WAAA,MAN0B,cAQhCE,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,EAAI,EAAIS,EAAItC,SAAS2C,KAAK,IAAED,YAC9CJ,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,CAxEgB,CAASwD,KA2EfC,cAAwBzD,GCtFvC0D,IAASC,OACPxB,cAACyB,IAAMC,WAAU,CAAA3B,SACfC,cAACnC,EAAe,MAElB8D,SAASC,eAAe,Q","file":"static/js/main.0c3ca5a6.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                {this.state.activeIndex != null ?\n                <>\n                    <strong key=\"index-label\">index:</strong> <span key=\"index\">{this.state.activeIndex} </span>\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 - 1 - 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\";\n\nReactDOM.render(\n  <React.StrictMode>\n    <HighlightedText />\n  </React.StrictMode>,\n  document.getElementById(\"root\")\n)\n"],"sourceRoot":""}

highlighted_text/build/static/js/main.243f5634.chunk.js

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

highlighted_text/build/static/js/main.243f5634.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","console","log","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":"2OAiBMA,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,YAC9BC,QAAQC,IAAIH,GAEZ,IAAII,EAAY,mBACZN,KAAKT,MAAMG,WACXY,GAAa,WAGjB,IAAMC,EAAU,WACZT,EAAKU,SAAS,CAAEd,UAAU,GAC9B,EAEA,OAAOe,eAAAC,WAAA,CAAAC,SAAA,CACHF,eAAA,OAAKH,UAAU,aAAYK,SAAA,CACvBF,eAAA,QAAMH,UAAWN,KAAKT,MAAMG,SAAW,GAAK,UAAUiB,SAAA,CAAiBC,cAAA,KAAGN,UAAU,eAAiB,MAA1C,aAC3DM,cAAA,UAAAD,SAA2B,WAAf,gBAEkB,MAA1BX,KAAKT,MAAMC,YACLoB,cAAA,QAAMN,UAAU,QAAOK,SAAeZ,EAAOC,KAAKT,MAAMC,cAA5B,UAC5BoB,cAAAF,WAAA,MANkB,cAShCE,cAAA,OAAKN,UAAWA,EAAWC,QAASA,EAAQI,SAEpCZ,EAAOc,KAAI,SAACC,EAAWC,GACnB,IAAIT,EAAY,QACZR,EAAKP,OACDO,EAAKP,MAAMC,aAAeuB,IAC1BT,GAAa,WAGrB,IAAMU,EAAQ,CACVC,gBACIf,EAAOa,GAAK,EAAC,qBAAAzB,OACcY,EAAOa,GAAE,0BAAAzB,QACRY,EAAOa,GAAE,MAS7C,OAAOH,cAAA,QAAcN,UAAWA,EAAWU,MAAOA,EAC9CE,YAPgB,WACXpB,EAAKP,MAAMG,UACZI,EAAKU,SAAS,CAAEhB,YAAauB,IAEjCjB,EAAKU,SAAS,CAAEf,WAAYsB,GAChC,EAE8BR,QAASA,EAAQI,SAAEG,GAD/BC,EAEtB,KAxByC,UA4BzD,GAAC,CAAAnB,IAAA,YAAAC,MAED,WACI,IAAME,EAASC,KAAKC,MAAMhB,KAAa,OACvC,IAAKe,KAAKT,OAAmC,MAA1BS,KAAKT,MAAMC,aAAuBQ,KAAKT,MAAMC,YAAc,EAC1E,OAAOO,EAAOc,KAAI,kBAAM,CAAC,IAE7B,IAAMM,EAAwBnB,KAAKC,MAAMhB,KAAa,OAEhD8B,EAAIf,KAAKT,MAAMC,YAAc,EAC7B4B,EAAKC,KAAKC,IAAID,KAAKE,IAAI,EAAGR,GAAII,EAAUJ,GAAG/B,QAC3CwC,EAAML,EAAUJ,GAAGU,MAAM,EAAGL,GAClCI,EAAIE,UACJ,IAAIC,EAAM,GAAArC,OAAAsC,YACH1C,MAAMmC,KAAKE,IAAI,EAAGR,EAAI,EAAIS,EAAIxC,SAAS6C,KAAK,IAAED,YAC9CJ,EAAIX,KAAI,SAACiB,GAAC,YAAUC,GAALD,GAAkBE,MAAMF,GAAK,EAAIA,CAAC,MAGxD,OADAH,EAAM,GAAArC,OAAAsC,YAAOD,GAAMC,YAAK1C,MAAMa,EAAOf,OAAS2C,EAAO3C,QAAQ6C,KAAK,IAEtE,KAACtD,CAAA,CA1EgB,CAAS0D,KA6EfC,cAAwB3D,GC1FvC4D,IAASC,OACPxB,cAACyB,IAAMC,WAAU,CAAA3B,SACfC,cAACrC,EAAe,MAElBgE,SAASC,eAAe,Q","file":"static/js/main.243f5634.chunk.js","sourcesContent":["import {\n    Streamlit,\n    StreamlitComponentBase,\n    withStreamlitConnection,\n} from \"streamlit-component-lib\";\nimport React, { ReactNode } from \"react\";\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        console.log(scores);\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                <strong key=\"target-label\">target:</strong>\n                {\n                    this.state.activeIndex != null\n                        ? <span className=\"token\" key=\"target\">{tokens[this.state.activeIndex]}</span>\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(255, 32, 32, ${scores[i]})`\n                                    : `rgba(32, 255, 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 - 1 - 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\";\n\nReactDOM.render(\n  <React.StrictMode>\n    <HighlightedText />\n  </React.StrictMode>,\n  document.getElementById(\"root\")\n)\n"],"sourceRoot":""}

highlighted_text/public/index.css

@@ -6,14 +6,11 @@ body {
     line-height: 1.4;
 }
 
-.status-bar {
-    display: none;
-}
-
 .highlighted-text {
     border: 1px solid #d2d2d2;
-    border-radius: 5px 5px 5px 5px;
-    padding: 2px;
+    border-radius: 0 0 5px 5px;
+    background-color: white;
+    padding: 4px;
     cursor: pointer;
 
     .token.active {
@@ -33,4 +30,21 @@ body {
             }
         }
     }
-}
+}
+
+.status-bar {
+    min-height: 1.4em;
+    border: 1px solid #d2d2d2;
+    border-bottom: none;
+    border-radius: 5px 5px 0 0;
+    padding: 3px 4px;
+
+    .token {
+        border: 1px solid #aaa;
+        border-radius: 2px;
+        background: #f5f5f5;
+        padding: 1px;
+        margin: 1px;
+        margin-right: 2px;
+    }
+}

highlighted_text/public/index.html

@@ -4,7 +4,6 @@
     <title>Highlighted text component</title>
     <meta charset="UTF-8" />
     <meta name="viewport" content="width=device-width, initial-scale=1" />
-    <meta name="theme-color" content="#fff" />
     <meta name="description" content="Highlighted text" />
     <link rel="stylesheet" href="index.css" />
   </head>

highlighted_text/src/HighlightedText.tsx

@@ -19,7 +19,6 @@ class HighlightedText extends StreamlitComponentBase<HighlightedTextState> {
     render() {
         const tokens: string[] = this.props.args["tokens"];
         const scores: number[] = this.getScores();
-        console.log(scores);
 
         let className = "highlighted-text";
         if (this.state.isFrozen) {
@@ -33,11 +32,12 @@ class HighlightedText extends StreamlitComponentBase<HighlightedTextState> {
         return <>
             <div className="status-bar" key="status-bar">
                 <span className={this.state.isFrozen ? "" : " d-none"} key="lock-icon"><i className="fa fa-lock"></i> </span>
-                <strong key="target-label">target:</strong>
                 {
-                    this.state.activeIndex != null
-                        ? <span className="token" key="target">{tokens[this.state.activeIndex]}</span>
-                        : <></>
+                    this.state.activeIndex != null ?
+                    <>
+                        <strong key="index-label">index:</strong> <span key="index">{this.state.activeIndex} </span>
+                    </>
+                    : <></>
                 }
             </div>
             <div className={className} onClick={onClick} key="text">
@@ -52,8 +52,8 @@ class HighlightedText extends StreamlitComponentBase<HighlightedTextState> {
                         const style = {
                             backgroundColor:
                                 scores[i] > 0
-                                    ? `rgba(255, 32, 32, ${scores[i]})`
-                                    : `rgba(32, 255, 32, ${-scores[i]})`
+                                    ? `rgba(32, 255, 32, ${scores[i]})`
+                                    : `rgba(255, 32, 32, ${-scores[i]})`
                         };
 
                         const onMouseOver = () => {