feat: Initial commit

requirements.txt

@@ -1,3 +1,4 @@
 altair
 pandas
-streamlit
+streamlit
+streamlit_vertical_slider

src/streamlit_app.py

@@ -1,40 +1,122 @@
 import altair as alt
-import numpy as np
 import pandas as pd
 import streamlit as st
+import streamlit_vertical_slider as svs
+import torch
+from streamlit_vertical_slider import vertical_slider
 
-"""
-# Welcome to Streamlit!
+st.title("Number Token Loss - Demo")
 
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
+st.markdown("""
+Adjust the sliders to set a predicted probability for each token (0-9 and "Text"). 
+The sliders are vertical and compact. The app normalizes the slider values 
+to form a valid probability distribution, visualizes it, and computes the corresponding 
+Cross Entropy, NTL-MSE, and NTL-WAS losses.
+""")
 
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
+# Vertical sliders for predicted probabilities of tokens 0-9 and "Text"
+st.markdown("#### Predicted Token Probabilities")
+cols = st.columns(11)
+prob_values = []
+for i, col in enumerate(cols):
+    label = f"Token {i}" if i < 10 else "Text"
+    with col:
+        val = svs.vertical_slider(
+            label=label,
+            min_value=0.0,
+            max_value=1.0,
+            step=0.1,
+            height=50,
+            key=f"slider_{i}",
+            slider_color="green",
+            track_color="lightgray",
+            thumb_color="black",
+        )
+        prob_values.append(val)
 
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
+# Normalize the probabilities to sum to 1
+total = sum(prob_values)
+probs = (
+    torch.ones(11) / 11.0
+    if total == 0
+    else torch.tensor([v / total for v in prob_values])
+)
 
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
+# Token labels
+options = [str(i) for i in range(10)] + ["Text"]
 
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
+# Ground truth token selection
+gt_choice = st.selectbox("Ground Truth Token", options=options, index=0)
+if gt_choice == "Text":
+    gt_index = 10
+    gt_numeric = None
+else:
+    gt_index = int(gt_choice)
+    gt_numeric = gt_index
 
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
+# Visualize the input distribution with highlighted ground truth bar
+st.markdown("#### Input Probability Distribution")
+df_dist = pd.DataFrame({"token": options, "probability": probs.numpy()})
+chart = (
+    alt.Chart(df_dist)
+    .mark_bar()
     .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))
+        x=alt.X("token:N", title="Token"),
+        y=alt.Y("probability:Q", title="Probability", scale=alt.Scale(domain=[0, 1])),
+        color=alt.condition(
+            alt.datum.token == gt_choice,
+            alt.value("green"),  # Highlight ground truth token
+            alt.value("steelblue"),  # Other tokens
+        ),
+    )
+    .properties(height=300)
+)
+st.altair_chart(chart, use_container_width=True)
+
+# Compute Cross Entropy loss: -log(predicted probability of the ground truth)
+ce_loss = -torch.log(torch.clamp(probs[gt_index], min=1e-9))
+
+# Compute NTL-MSE loss
+if gt_numeric is None:
+    ntl_mse_loss = torch.tensor(0.0)
+else:
+    numeric_probs = probs[:10]
+    values = torch.arange(0, 10, dtype=torch.float32)
+    pred_value = torch.sum(numeric_probs * values)
+    ntl_mse_loss = (pred_value - float(gt_numeric)) ** 2
+
+# Compute NTL-WAS loss
+if gt_numeric is None:
+    ntl_was_loss = torch.tensor(0.0)
+else:
+    numeric_probs = probs[:10]
+    values = torch.arange(0, 10, dtype=torch.float32)
+    abs_diff = torch.abs(values - float(gt_numeric))
+    ntl_was_loss = torch.sum(numeric_probs * abs_diff)
+
+# Convert losses to Python floats and round to 3 decimals
+ce_val = round(ce_loss.item(), 3)
+mse_val = round(ntl_mse_loss.item(), 3)
+was_val = round(ntl_was_loss.item(), 3)
+
+# Display numeric values of the losses
+st.subheader("Loss Values")
+st.write(f"**Cross Entropy:** {ce_val:.3f}")
+st.write(f"**NTL-MSE:** {mse_val:.3f}")
+st.write(f"**NTL-WAS:** {was_val:.3f}")
+
+# Bar chart comparison of the three losses
+st.subheader("Loss Comparison Chart")
+loss_df = pd.DataFrame(
+    {
+        "Loss": ["Cross Entropy", "NTL-MSE", "NTL-WAS"],
+        "Value": [ce_val, mse_val, was_val],
+    }
+).set_index("Loss")
+st.bar_chart(loss_df)
+
+# References / resources section with links
+st.markdown("### Resources")
+st.markdown(
+    "- **Paper:** [Regress, Don't Guess – A Regression-like Loss on Number Tokens for Language Models](https://arxiv.org/abs/2411.02083)  \n- **Code:** [tum-ai/number-token-loss (GitHub)](https://github.com/tum-ai/number-token-loss)"
+)