Cache the results

app.py

@@ -62,28 +62,40 @@ 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))
-tokens = ids_to_readable_tokens(tokenizer, 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))
+
+    inputs_sliding = get_windows_batched(
+        inputs,
+        window_len=window_len,
+        pad_id=tokenizer.eos_token_id
+    )
+    with torch.inference_mode():
+        logits = model(**inputs_sliding.convert_to_tensors("pt")).logits.to(torch.float16)
+        logits = logits.permute(1, 0, 2)
+        logits = F.pad(logits, (0, 0, 0, window_len, 0, 0), value=torch.nan)
+        logits = logits.view(-1, logits.shape[-1])[:-window_len]
+        logits = logits.view(window_len, len(input_ids) + window_len - 2, logits.shape[-1])
 
-inputs_sliding = get_windows_batched(
-    inputs,
-    window_len=window_len,
-    pad_id=tokenizer.eos_token_id
+    scores = logits.to(torch.float32).log_softmax(dim=-1)
+    scores = scores[:, torch.arange(len(input_ids[1:])), input_ids[1:]]
+    scores = scores.diff(dim=0).transpose(0, 1)
+    scores = scores.nan_to_num()
+    scores /= scores.abs().max(dim=1, keepdim=True).values + 1e-9
+    scores = scores.to(torch.float16)
+
+    return input_ids, scores
+
+input_ids, scores = run_context_length_probing(
+    model_name=model_name,
+    text=text,
+    window_len=window_len
 )
-with torch.inference_mode():
-    logits = model(**inputs_sliding.convert_to_tensors("pt")).logits.to(torch.float16)
-    logits = logits.permute(1, 0, 2)
-    logits = F.pad(logits, (0, 0, 0, window_len, 0, 0), value=torch.nan)
-    logits = logits.view(-1, logits.shape[-1])[:-window_len]
-    logits = logits.view(window_len, len(input_ids) + window_len - 2, logits.shape[-1])
-
-scores = logits.to(torch.float32).log_softmax(dim=-1)
-scores = scores[:, torch.arange(len(input_ids[1:])), input_ids[1:]]
-scores = scores.diff(dim=0).transpose(0, 1)
-scores = scores.nan_to_num()
-scores /= scores.abs().max(dim=1, keepdim=True).values + 1e-9
-scores = scores.to(torch.float16)
+tokens = ids_to_readable_tokens(tokenizer, input_ids)
 
 highlighted_text_component(tokens=tokens, scores=scores.tolist())