Better progress indication

app.py

@@ -56,8 +56,9 @@ if metric_name == "KL divergence":
     st.error("KL divergence is not supported yet. Stay tuned!", icon="😭")
     st.stop()
 
-tokenizer = st.cache_resource(AutoTokenizer.from_pretrained, show_spinner=False)(model_name)
-model = st.cache_resource(AutoModelForCausalLM.from_pretrained, show_spinner=False)(model_name)
+with st.spinner("Loading model…"):
+    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"]
@@ -80,29 +81,29 @@ def run_context_length_probing(model_name, text, window_len):
     ).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)
-    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)
-
-    pbar.progress(1., "Done!")
+    with st.spinner("Running model…"):
+        batch_size = 8
+        num_items = len(inputs_sliding["input_ids"])
+        pbar = st.progress(0)
+        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 = torch.cat(logits, dim=0)
+        pbar.empty()
+
+    with st.spinner("Computing scores…"):
+        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)
 
     return scores