Add TinyStories models

README.md

@@ -12,4 +12,6 @@ models:
 - distilgpt2
 - gpt2
 - EleutherAI/gpt-neo-125m
+- roneneldan/TinyStories-8M
+- roneneldan/TinyStories-33M
 ---

app.py

@@ -99,7 +99,16 @@ st.caption(
     "and visualize the contributions of different contexts to each generated token."
 )
 
-model_name = st.selectbox("Model", ["distilgpt2", "gpt2", "EleutherAI/gpt-neo-125m"])
+model_name = st.selectbox(
+    "Model",
+    [
+        "distilgpt2",
+        "gpt2",
+        "EleutherAI/gpt-neo-125m",
+        "roneneldan/TinyStories-8M",
+        "roneneldan/TinyStories-33M",
+    ]
+)
 metric_name = st.radio(
     "Metric", (["KL divergence"] if not generation_mode else []) + ["NLL loss"], index=0, horizontal=True
 )