Better caching

app.py

@@ -62,13 +62,14 @@ if metric_name == "KL divergence":
 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))
+
 @st.cache_data(show_spinner=False)
 def run_context_length_probing(model_name, text, window_len):
     assert model.name_or_path == model_name
-
-    inputs = tokenizer([text])
-    [input_ids] = inputs["input_ids"]
-    window_len = min(window_len, len(input_ids))
+    del text  # needed as a cache key but for the computation we access inputs directly
 
     inputs_sliding = get_windows_batched(
         inputs,
@@ -89,9 +90,9 @@ def run_context_length_probing(model_name, text, window_len):
     scores /= scores.abs().max(dim=1, keepdim=True).values + 1e-9
     scores = scores.to(torch.float16)
 
-    return input_ids, scores
+    return scores
 
-input_ids, scores = run_context_length_probing(
+scores = run_context_length_probing(
     model_name=model_name,
     text=text,
     window_len=window_len