Add more efficient processing

app.py

@@ -120,9 +120,7 @@ def process_texts(student_model_id, teacher_model_id, dataset_id, progress=gr.Pr
     ds = load_dataset(dataset_id, split="train")
     rows = ds.select(range(min(10, len(ds))))
 
-    html_blocks = []
-    for row_idx, row in enumerate(rows):
-        progress((row_idx + 1) / 12, desc=f"Processing text {row_idx + 1}/10...")
+    def make_html_block(row, idx):
         text = "".join(msg["content"] for msg in row["messages"])
 
         s_ids = student_tokenizer.encode(text, add_special_tokens=False)
@@ -133,13 +131,19 @@ def process_texts(student_model_id, teacher_model_id, dataset_id, progress=gr.Pr
         )
 
         highlighted = highlight_groups(student_tokenizer, s_ids, s_groups, t_groups)
-        html_blocks.append(
-            f'<div style="border:1px solid #ccc; padding:10px; margin:10px 0; '
-            f'border-radius:5px; white-space:pre-wrap; font-family:monospace; font-size:13px;">'
-            f"<strong>Text {row_idx + 1}</strong> "
-            f"(student tokens: {len(s_ids)}, teacher tokens: {len(t_ids)})<br><br>"
-            f"{highlighted}</div>"
-        )
+        return {
+            "html_block": (
+                f'<div style="border:1px solid #ccc; padding:10px; margin:10px 0; '
+                f'border-radius:5px; white-space:pre-wrap; font-family:monospace; font-size:13px;">'
+                f"<strong>Text {idx + 1}</strong> "
+                f"(student tokens: {len(s_ids)}, teacher tokens: {len(t_ids)})<br><br>"
+                f"{highlighted}</div>"
+            )
+        }
+
+    progress(0.2, desc="Processing texts...")
+    rows = rows.map(make_html_block, num_proc=4, with_indices=True)
+    html_blocks = rows["html_block"]
 
     progress(1, desc="Done!")