app.py

import numpy as np
import pandas as pd
import random
import solara
import torch
import torch.nn.functional as F
import ipyvue
import reacton
from solara.alias import rv as v
from typing import Any, Callable, Optional, TypeVar, Union, cast, overload
from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig

config = AutoConfig.from_pretrained(
    "replit/replit-code-v1_5-3b",
    trust_remote_code=True
)

tokenizer = AutoTokenizer.from_pretrained('replit/replit-code-v1_5-3b', trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained('replit/replit-code-v1_5-3b', config=config, trust_remote_code=True)

def use_change(el: reacton.core.Element, on_value: Callable[[Any], Any], enabled=True):
    """Trigger a callback when a blur events occurs or the enter key is pressed."""
    on_value_ref = solara.use_ref(on_value)
    on_value_ref.current = on_value

    def add_events():
        def on_change(widget, event, data):
            if enabled:
                on_value_ref.current(widget.v_model)

        widget = cast(ipyvue.VueWidget, solara.get_widget(el))
        if enabled:
            widget.on_event("blur", on_change)
            widget.on_event("keyup.enter", on_change)

        def cleanup():
            if enabled:
                widget.on_event("blur", on_change, remove=True)
                widget.on_event("keyup.enter", on_change, remove=True)

        return cleanup

    solara.use_effect(add_events, [enabled])


@solara.component
def InputTextarea(
    label: str,
    value: Union[str, solara.Reactive[str]] = "",
    on_value: Callable[[str], None] = None,
    disabled: bool = False,
    password: bool = False,
    continuous_update: bool = False,
    error: Union[bool, str] = False,
    message: Optional[str] = None,
):
    reactive_value = solara.use_reactive(value, on_value)
    del value, on_value

    def set_value_cast(value):
        reactive_value.value = str(value)

    def on_v_model(value):
        if continuous_update:
            set_value_cast(value)

    messages = []
    if error and isinstance(error, str):
        messages.append(error)
    elif message:
        messages.append(message)
    text_area = v.Textarea(
        v_model=reactive_value.value,
        on_v_model=on_v_model,
        label=label,
        disabled=disabled,
        type="password" if password else None,
        error=bool(error),
        messages=messages,
        solo=True,
        hide_details=True,
        outlined=True,
        rows=1,
        auto_grow=True,
    )
    use_change(text_area, set_value_cast, enabled=not continuous_update)
    return text_area

@solara.component
def my_component(tokens, i, color, df):
  text = tokenizer.decode(tokens[0][i+1])
  color = solara.use_reactive(f"{color}")
  text_element = solara.Text(f"{text}", classes=[f"{color.value}"])
  with solara.lab.ClickMenu(activator=text_element):
    with solara.Column(gap="0px"):
      def replace_token(text=text):
        color.set("mystronggreen")
        text1.value = f"{tokenizer.decode(tokens[0][1:i+1])}"+f"{df.iloc[1,1]}"+f"{tokenizer.decode(tokens[0][i+2:])}"
      solara.Button(f"Replace "+ f"'{tokenizer.decode(tokens[0][i+1])}'".replace(" ", "␣")+" by "+f"'{df.iloc[1,1]}'".replace(" ", "␣"), on_click=replace_token, text=True, classes=["mybuttonclass"])
      def add_token(text=text):
        color.set("mystronggreen")
        text1.value = f"{tokenizer.decode(tokens[0][1:i+1])}"+f"{df.iloc[1,1]}"+f"{tokenizer.decode(tokens[0][i+1:])}"
      solara.Button(f"Add "+f"'{df.iloc[1,1]}'".replace(" ", "␣"), on_click=add_token, text=True, classes=["mybuttonclass"])
      def delete_token(text=text):
        color.set("mystronggreen")
        text1.value = f"{tokenizer.decode(tokens[0][1:i+1])}"+f"{tokenizer.decode(tokens[0][i+2:])}"
      solara.Button(f"Delete "+f"'{tokenizer.decode(tokens[0][i+1])}'".replace(" ", "␣"), on_click=delete_token, text=True, classes=["mybuttonclass"])
      def ignore_token(text=text):
        color.set("mystronggreen")
      solara.Button("Ignore", on_click=ignore_token, text=True, classes=["mybuttonclass"])

text1 = solara.reactive("""def HelloWorld():\n  print("Hello World)""")
@solara.component
def Page():
  with solara.Column(margin="10"):
    solara.Markdown("#Code Perplexity")
    solara.Markdown("This is an educational tool where, for any given passage of code, it augments the original code with highlights and annotations that indicate how 'surprising' each token is to the model, as well as which other tokens the model deemed most likely to occur in its place.")
    css = """
    .mybuttonclass{
      text-transform: none !important;
    }
    .mystronggreen{
      background-color:#99ff99;
      color:black!important;
      padding:0px;
      white-space-collapse:preserve;
    }
    .mygreen{
      background-color:#ccffcc;
      color:black!important;
      white-space-collapse:preserve;
    }
    .myyellow{
      background-color: #ffff99;
      color:black!important;
      white-space-collapse:preserve;
    }
    .myorange{
      background-color: #ffe6cc;
      color:black!important;
      white-space-collapse:preserve;
    }
    .myred{
      background-color:#ffcab0;
      color:black!important;
      white-space-collapse:preserve;
    }
    """
    InputTextarea("Enter text and press enter when you're done:", value=text1, continuous_update=True)
    if text1.value != "":
      with solara.Column():
        with solara.Row(gap="0px", justify="left"):
          tokens = tokenizer.encode(text1.value, return_tensors="pt")
          tokens = torch.concat((torch.tensor([tokenizer.eos_token_id]), tokens[0])).reshape(1,-1)
          full_list = []
          partial_list = []
          for token in tokens[0]:
            if token != 216:
              partial_list.append(token)
            else:
              partial_list.append(torch.tensor(216))
              full_list.append(partial_list)
              partial_list = []
          if len(partial_list) != 0:
            full_list.append(partial_list)
          tokens = torch.cat((torch.tensor([tokenizer.eos_token_id]), tokens[0])).reshape(1,-1)
          # tokens = tokens[0].reshape(1,-1)
          i = 0
          for j in range(len(full_list)):
            with solara.Column():
              with solara.Div(style="display: inline;"):
                for k in range(len(full_list[j])):
                  outputs = model.generate(tokens[0][:i+1].reshape(1,-1), max_new_tokens=1, output_scores=True, return_dict_in_generate=True, eos_token_id=tokenizer.eos_token_id, pad_token_id=tokenizer.eos_token_id)
                  scores = F.softmax(outputs.scores[0], dim=-1)
                  top_10 = torch.topk(scores, 10)
                  df = pd.DataFrame()
                  a = scores[0][tokens[0][i+1]]
                  b = top_10.values
                  df["probs"] = list(np.concatenate([a.reshape(-1,1).numpy()[0], b[0].numpy()]))
                  diff = 100*(df["probs"].iloc[0]-df["probs"].iloc[1])
                  if np.abs(diff)<1:
                    color = "mystronggreen"
                  elif np.abs(diff)<10:
                    color = "mygreen"
                  elif np.abs(diff)<20:
                    color = "myyellow"
                  elif np.abs(diff)<30:
                    color = "myorange"
                  else:
                    color = "myred"
                  df["probs"] = [f"{value:.2%}" for value in df["probs"].values]
                  aux = [tokenizer.decode(tokens[0][i+1])] + [tokenizer.decode(top_10.indices[0][i]) for i in range(10)]
                  df["predicted next token"] = aux
                  solara_df = solara.DataFrame(df, items_per_page=10)
                  with solara.Tooltip(solara_df, color="white"):
                    solara.Style(css)
                    if full_list[j][k] == 216:
                      solara.Text("↵", classes=[f"{color}"])
                    elif full_list[j][k] == 0:
                      solara.Text("")
                    else:
                      solara.Text(f"{tokenizer.decode(full_list[j][k])}", classes=[f"{color}"])
                  i+=1