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 tokenizer = AutoTokenizer.from_pretrained('gpt2', padding_side='left') model = AutoModelForCausalLM.from_pretrained('gpt2') 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 css = """ .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; } """ def get_color(diff): 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" return color text1 = solara.reactive("""If I start counting to ten... one two three four five six seven eight nine ten ... you can see that as the sequence progresses, the LM gets better at predicting the sequence and by the end, it's 100% correct about how things are going to continue. The first time I write this sentence, the model is quite confused about what token is about to come next, especially if I throw in weird words like pumpkin, clown, tweets, alpha, teddy bear. But the second time? It expects almost every word that appears: The first time I write this sentence, the model is quite confused about what token is about to come next, especially if I throw in weird words like pumpkin, clown, tweets, alpha, teddy bear.""") @solara.component def Page(): with solara.Column(margin="10"): solara.Markdown("#Perplexity") solara.Markdown("This is an educational tool. For any given passage of text, this tool augments the original text 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 when you hover over each token.") InputTextarea("Enter text:", value=text1, continuous_update=True) if text1.value != "": tokens = tokenizer.encode(text1.value, return_tensors="pt") full_list = [] partial_list = [] for token in tokens[0]: if token != 198: partial_list.append(token) else: partial_list.append(torch.tensor(198)) 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) 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, 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 c = torch.concat([a.reshape(-1,1)[0], b[0]]) df["probs"] = list([c[i].cpu().item() for i in range(len(c))]) diff = 100*(df["probs"].iloc[0]-df["probs"].iloc[1]) color = get_color(diff) df["probs"] = [f"{value:.2%}" for value in df["probs"].values] aux = [tokenizer.decode(tokens[0][i+1]).replace(" ", "␣")] + [tokenizer.decode(top_10.indices[0][i]).replace(" ", "␣") for i in range(10)] df["predicted next token"] = aux solara_df = solara.DataFrame(df, items_per_page=11) with solara.Tooltip(solara_df, color="white"): solara.Style(css) if full_list[j][k] == 198: solara.Text("↵", classes=[f"{color}"]) else: solara.Text(f"{tokenizer.decode(full_list[j][k])}", classes=[f"{color}"]) i+=1 Page()