Add progress bar, batched inference

app.py

@@ -77,9 +77,18 @@ def run_context_length_probing(model_name, text, window_len):
         inputs,
         window_len=window_len,
         pad_id=tokenizer.eos_token_id
-    )
-
-    logits = model(**inputs_sliding.convert_to_tensors("pt")).logits.to(torch.float16)
+    ).convert_to_tensors("pt")
+
+    logits = []
+    pbar = st.progress(0.)
+    batch_size = 8
+    num_items = len(inputs_sliding["input_ids"])
+    for i in range(0, num_items, batch_size):
+        pbar.progress(i / num_items * 0.9, f"Running model… ({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))
+    pbar.progress(0.9, "Computing scores…")
+    logits = torch.cat(logits, dim=0)
 
     logits = logits.permute(1, 0, 2)
     logits = F.pad(logits, (0, 0, 0, window_len, 0, 0), value=torch.nan)
@@ -93,6 +102,8 @@ 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)
 
+    pbar.progress(1., "Done!")
+
     return scores
 
 scores = run_context_length_probing(