Cache models, add gpt-neo-125m

app.py

@@ -44,12 +44,9 @@ def ids_to_readable_tokens(tokenizer, ids, strip_whitespace=False):
 st.header("Context length probing")
 
 with st.form("form"):
-    model_name = st.selectbox("Model", ["distilgpt2", "gpt2"])
+    model_name = st.selectbox("Model", ["distilgpt2", "gpt2", "EleutherAI/gpt-neo-125m"])
     metric_name = st.selectbox("Metric", ["Cross entropy"])
 
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(model_name)
-
     window_len = st.select_slider("Window size", options=[8, 16, 32, 64, 128, 256, 512, 1024], value=512)
     text = st.text_area(
         "Input text",
@@ -58,6 +55,9 @@ with st.form("form"):
 
     st.form_submit_button("Submit")
 
+tokenizer = st.cache_resource(AutoTokenizer.from_pretrained, show_spinner=False)(model_name)
+model = st.cache_resource(AutoModelForCausalLM.from_pretrained, show_spinner=False)(model_name)
+
 inputs = tokenizer([text])
 [input_ids] = inputs["input_ids"]
 window_len = min(window_len, len(input_ids))