highlighted_text/src/HighlightedText.tsx

import {
    StreamlitComponentBase,
    withStreamlitConnection,
} from "streamlit-component-lib";

type HighlightedTextState = {
    activeIndex: number | null,
    hoverIndex: number | null,
    isFrozen: boolean
};

/**
 * This is a React-based component template. The `render()` function is called
 * automatically when your component should be re-rendered.
 */
class HighlightedText extends StreamlitComponentBase<HighlightedTextState> {
    public state = {activeIndex: null, hoverIndex: null, isFrozen: false};

    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) {
            className += " frozen";
        }

        const onClick = () => {
            this.setState({ isFrozen: false });
        };

        return <div className="container">
            <div className="status-bar" key="status-bar">
                <span className={this.state.isFrozen ? "" : " d-none"} key="lock-icon"><i className="fa fa-lock"></i> </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">
                {
                    tokens.map((t: string, i: number) => {
                        let className = "token";
                        if (this.state) {
                            if (this.state.activeIndex == i) {
                                className += " active";
                            }
                        }
                        if (i < prefixLength) {
                            className += " prefix";
                        }
                        const style = {
                            backgroundColor:
                                scores[i] > 0
                                    ? `rgba(32, 255, 32, ${scores[i]})`
                                    : `rgba(255, 32, 32, ${-scores[i]})`
                        };

                        const onMouseOver = () => {
                            if (!this.state.isFrozen) {
                                this.setState({ activeIndex: i });
                            }
                            this.setState({ hoverIndex: i });
                        };
                        return <span key={i} className={className} style={style}
                            onMouseOver={onMouseOver} onClick={onClick}>{t}</span>;
                    })
                }
            </div>
        </div>;
    }

    private getScores() {
        const tokens = this.props.args["tokens"];
        if (!this.state || this.state.activeIndex == null || this.state.activeIndex < 1) {
            return tokens.map(() => 0);
        }
        const allScores: number[][] = this.props.args["scores"];

        const i = this.state.activeIndex - 1;
        const hi = Math.min(Math.max(0, i + 1), allScores[i].length);
        const row = allScores[i].slice(0, hi);
        row.reverse();
        let result = [
            ...Array(Math.max(0, i + 1 - row.length)).fill(0),
            ...row.map((x) => x == undefined || isNaN(x) ? 0 : x)
        ];
        result = [...result, ...Array(tokens.length - result.length).fill(0)];
        return result;
    }
}

export default withStreamlitConnection(HighlightedText);