Cache individual batches

app.py

@@ -70,14 +70,19 @@ if len(input_ids) < 2:
 
 @st.cache_data(show_spinner=False)
 @torch.inference_mode()
-def run_context_length_probing(model_name, text, window_len):
-    assert model.name_or_path == model_name
-    del text  # needed as a cache key but for the computation we access inputs directly
+def get_logits(_model, _inputs, cache_key):
+    del cache_key
+    return _model(**_inputs).logits.to(torch.float16)
+
+@st.cache_data(show_spinner=False)
+@torch.inference_mode()
+def run_context_length_probing(_model, _tokenizer, _inputs, window_len, cache_key):
+    del cache_key
 
     inputs_sliding = get_windows_batched(
-        inputs,
+        _inputs,
         window_len=window_len,
-        pad_id=tokenizer.eos_token_id
+        pad_id=_tokenizer.eos_token_id
     ).convert_to_tensors("pt")
 
     logits = []
@@ -88,7 +93,13 @@ def run_context_length_probing(model_name, text, window_len):
         for i in range(0, num_items, batch_size):
             pbar.progress(i / num_items, f"{i}/{num_items}")
             batch = {k: v[i:i + batch_size] for k, v in inputs_sliding.items()}
-            logits.append(model(**batch).logits.to(torch.float16))
+            logits.append(
+                get_logits(
+                    _model,
+                    batch,
+                    cache_key=(model_name, batch["input_ids"].cpu().numpy().tobytes())
+                )
+            )
         logits = torch.cat(logits, dim=0)
         pbar.empty()
 
@@ -108,9 +119,11 @@ def run_context_length_probing(model_name, text, window_len):
     return scores
 
 scores = run_context_length_probing(
-    model_name=model_name,
-    text=text,
-    window_len=window_len
+    _model=model,
+    _tokenizer=tokenizer,
+    _inputs=inputs,
+    window_len=window_len,
+    cache_key=(model_name, text),
 )
 tokens = ids_to_readable_tokens(tokenizer, input_ids)