Update app.py

app.py

@@ -1,40 +1,96 @@
-import solara as sol
-from transformers import AutoModelForCausalLM, AutoTokenizer
+Hugging Face's logo
+Hugging Face
+Search models, datasets, users...
+Models
+Datasets
+Spaces
+Posts
+Docs
+Pricing
+
+
+
+Spaces:
+
+alonsosilva
+/
+NextTokenPrediction
+
+
+like
+3
+App
+Files
+Community
+NextTokenPrediction
+/
+app.py
+
+alonsosilva's picture
+alonsosilva
+Change reactive text
+a4869ab
+7 months ago
+raw
+history
+blame
+contribute
+delete
+No virus
+2.63 kB
+import solara
+import random
 import torch
+import torch.nn.functional as F
+import pandas as pd
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+tokenizer = AutoTokenizer.from_pretrained('gpt2')
+model = AutoModelForCausalLM.from_pretrained('gpt2')
+text1 = solara.reactive("Never gonna give you up, never gonna let you")
+@solara.component
+def Page():
+  with solara.Column(margin=10):
+    solara.Markdown("#Next token prediction visualization")
+    solara.Markdown("I built this tool to help me understand autoregressive language models. For any given text, it gives the top 10 candidates to be the next token with their respective probabilities. The language model I'm using is the smallest version of GPT-2, with 124M parameters.")
+    def on_action_cell(column, row_index):
+      text1.value += tokenizer.decode(top_10.indices[0][row_index])
+    cell_actions = [solara.CellAction(icon="mdi-thumb-up", name="Select", on_click=on_action_cell)]
+    solara.InputText("Enter text:", value=text1, continuous_update=True)
+    if text1.value != "":
+      tokens = tokenizer.encode(text1.value, return_tensors="pt")
+      spans1 = ""
+      spans2 = ""
+      for i, token in enumerate(tokens[0]):
+        random.seed(i)
+        random_color = ''.join([random.choice('0123456789ABCDEF') for k in range(6)])
+        spans1 += " " + f"<span style='font-family: helvetica; color: #{random_color}'>{token}</span>"
+        spans2 += " " + f"""<span style="
+            padding: 6px;
+            border-right: 3px solid white;
+            line-height: 3em;
+            font-family: courier;
+            background-color: #{random_color};
+            color: white;
+            position: relative;
+          "><span style="
+          position: absolute;
+          top: 5.5ch;
+          line-height: 1em;
+          left: -0.5px;
+          font-size: 0.45em"> {token}</span>{tokenizer.decode([token])}</span>"""
+      solara.Markdown(f'{spans2}')
+      solara.Markdown(f'{spans1}')
+      outputs = model.generate(tokens, 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()
+      df["probs"] = top_10.values[0]
+      df["probs"] = [f"{value:.2%}" for value in df["probs"].values]
+      df["next token ID"] = [top_10.indices[0][i].numpy() for i in range(10)]
+      df["predicted next token"] = [tokenizer.decode(top_10.indices[0][i]) for i in range(10)]
+      solara.Markdown("###Prediction")
+      solara.DataFrame(df, items_per_page=10, cell_actions=cell_actions)
+Page()
 
-model_name = "datificate/gpt2-small-spanish"
-model = AutoModelForCausalLM.from_pretrained(model_name)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-
-def predict_next_token(text):
-    inputs = tokenizer(text, return_tensors="pt")
-    outputs = model(**inputs)
-    next_token_logits = outputs.logits[:, -1, :]
-    next_token_probs = torch.softmax(next_token_logits, dim=-1)
-    top_k_probs, top_k_indices = torch.topk(next_token_probs, 10)
-    top_k_tokens = tokenizer.convert_ids_to_tokens(top_k_indices[0])
-    return list(zip(top_k_tokens, top_k_probs[0].tolist()))
-
-@sol.component
-def NextTokenPredictionApp():
-    text = sol.reactive("")
-    predictions = sol.reactive([])
-
-    def on_text_change(new_text):
-        text.set(new_text)
-        preds = predict_next_token(new_text)
-        predictions.set(preds)
-
-    sol.Markdown("# Predicción del Próximo Token")
-    sol.Markdown("Ingrese un texto en español y vea las predicciones para el próximo token.")
-    
-    sol.InputText(value=text.value, on_change=on_text_change, placeholder="Escribe algo en español...", fullwidth=True)
-    sol.Button("Predecir", on_click=lambda: on_text_change(text.value))
-
-    if predictions.value:
-        sol.Markdown("## Predicciones de tokens:")
-        for token, prob in predictions.value:
-            sol.Markdown(f"- **{token}**: {prob:.4f}")
-
-# Iniciar la aplicación en modo de desarrollo
-app = sol.App(NextTokenPredictionApp, title="Next Token Prediction App")
+