Upload 2 files

app.py

@@ -2891,8 +2891,7 @@ def _chat_stream(history, version, ckpt_label, mode_key, use_custom,
 
 def _compare_fn(prompt, selected_versions, mode_key, use_custom,
                 temperature, top_k, top_p, min_p, rep_penalty,
-                ngram_size, soft_cap, loop_pen, max_tokens, ctx_win, raw_mode,
-                progress=gr.Progress(track_tqdm=True)):
+                ngram_size, soft_cap, loop_pen, max_tokens, ctx_win, raw_mode):
     if use_custom:
         cfg = {
             "sft_mode": not raw_mode,
@@ -2906,16 +2905,39 @@ def _compare_fn(prompt, selected_versions, mode_key, use_custom,
         cfg = dict(MODES[mode_key])
 
     all_versions = _collection_versions()
-    results = {}
-    for version in progress.tqdm(selected_versions or [], desc="Running models"):
+    selected = set(selected_versions or [])
+    state = {v: ("⏳ Queued…" if v in selected else "") for v in all_versions}
+
+    def _emit():
+        return [state[v] for v in all_versions]
+
+    yield _emit()
+
+    for version in all_versions:
+        if version not in selected:
+            continue
         labels = _ckpt_labels(version)
         ckpt_label = labels[0] if labels else None
         if not ckpt_label:
-            results[version] = "[No checkpoint found]"
+            state[version] = "[No checkpoint found]"
+            yield _emit()
             continue
+
+        state[version] = "⏳ Loading…"
+        yield _emit()
+
         try:
             bundle = _load_bundle(version, ckpt_label)
-            out = generate(
+        except Exception as e:
+            state[version] = f"[Load error: {e}]"
+            yield _emit()
+            continue
+
+        state[version] = ""
+        yield _emit()
+
+        try:
+            for partial in generate_stream(
                 model=bundle["model"], tokenizer=bundle["tokenizer"],
                 prompt=prompt, device=str(bundle["device"]),
                 sft_mode=cfg["sft_mode"],
@@ -2927,14 +2949,12 @@ def _compare_fn(prompt, selected_versions, mode_key, use_custom,
                 loop_penalty=cfg["loop_penalty"],
                 max_new_tokens=cfg["max_new_tokens"],
                 context_window=cfg["context_window"],
-                stream=False,
-            )
-            results[version] = out
+            ):
+                state[version] = partial
+                yield _emit()
         except Exception as e:
-            results[version] = f"[Error: {e}]"
-
-    # Return one value per discovered version (empty string if not selected/run)
-    return [results.get(v, "") for v in all_versions]
+            state[version] = f"[Generation error: {e}]"
+            yield _emit()
 
 
 # ---- benchmark ----
@@ -3090,24 +3110,25 @@ with gr.Blocks(title="CompactAI Models") as demo:
 
     # ── Compare ───────────────────────────────────────────────────────────────
     with gr.Tab("Compare All Models"):
-        gr.Markdown("Run the same prompt on multiple models and compare side-by-side.")
+        gr.Markdown("Run the same prompt on every selected model. Outputs stream live one model at a time.")
         with gr.Row():
-            cmp_prompt = gr.Textbox(label="Prompt", placeholder="Enter a prompt…", lines=3, scale=3)
+            cmp_prompt = gr.Textbox(label="Prompt", placeholder="Enter a prompt…", lines=4, scale=3)
             with gr.Column(scale=1):
                 cmp_models = gr.CheckboxGroup(
-                    choices=_initial_versions, value=_initial_versions, label="Models"
+                    choices=_initial_versions, value=_initial_versions, label="Models to run"
                 )
                 cmp_mode = gr.Dropdown(
                     choices=_mode_keys, value="chat-coherent", label="Mode preset"
                 )
         cmp_use_custom, cmp_temp, cmp_topk, cmp_topp, cmp_minp, cmp_rep, cmp_ng, cmp_cap, cmp_lp, cmp_maxt, cmp_ctx, cmp_raw = _advanced_block()
-        cmp_run = gr.Button("Run comparison", variant="primary")
-
-        with gr.Row():
-            cmp_outputs = [
-                gr.Textbox(label=v, lines=8)
-                for v in _initial_versions
-            ]
+        cmp_run = gr.Button("▶ Run comparison", variant="primary")
+
+        # 2-column grid of output boxes
+        cmp_outputs = []
+        for row_start in range(0, len(_initial_versions), 2):
+            with gr.Row():
+                for v in _initial_versions[row_start:row_start + 2]:
+                    cmp_outputs.append(gr.Textbox(label=v, lines=10, interactive=False))
 
         cmp_run.click(
             _compare_fn,