Improve Colab runner progress visibility

eval/news_summarization/run_hf_transformers.py

@@ -29,7 +29,7 @@ from eval.news_summarization.run_news_summary_pilot import (  # noqa: E402
 
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(description="Run the news summarization pilot on Hugging Face/Colab with transformers.")
-    parser.add_argument("--model", required=True, help="Hub model id, e.g. Qwen/Qwen3.5-8B-Instruct-2507")
+    parser.add_argument("--model", required=True, help="Hub model id, e.g. Qwen/Qwen2.5-7B-Instruct")
     parser.add_argument("--dataset", default="bbc2024_qwen_reference")
     parser.add_argument("--prompt-style", default="simple", choices=["simple", "helpful", "detailed"])
     parser.add_argument("--limit", type=int, default=50)
@@ -45,7 +45,7 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--bertscore-model", default="roberta-large")
     parser.add_argument("--output")
     parser.add_argument("--resume", action="store_true")
-    parser.add_argument("--save-every", type=int, default=5)
+    parser.add_argument("--save-every", type=int, default=1)
     parser.add_argument("--verbose", action="store_true")
     return parser.parse_args()
 
@@ -54,24 +54,21 @@ def build_generator(args: argparse.Namespace):
     import torch
     from transformers import pipeline
 
-    torch_dtype = None
+    dtype = None
     if args.dtype != "auto":
-        torch_dtype = getattr(torch, args.dtype)
+        dtype = getattr(torch, args.dtype)
 
     kwargs: dict[str, object] = {
         "model": args.model,
         "device_map": args.device_map,
         "trust_remote_code": args.trust_remote_code,
     }
-    if torch_dtype is not None:
-        kwargs["torch_dtype"] = torch_dtype
+    if dtype is not None:
+        kwargs["dtype"] = dtype
     if args.attn_implementation:
         kwargs["model_kwargs"] = {"attn_implementation": args.attn_implementation}
 
     generator = pipeline("text-generation", **kwargs)
-    generation_config = getattr(generator.model, "generation_config", None)
-    if generation_config is not None and getattr(generation_config, "max_length", None) == 20:
-        generation_config.max_length = None
     return generator
 
 
@@ -113,6 +110,7 @@ def main() -> int:
         print(f"Resuming from {output_path} with {len(rows)} completed cases.")
 
     pending_cases = [case for case in cases if case.case_id not in completed_case_ids]
+    total_cases = len(cases)
     for index, case in enumerate(pending_cases, start=len(rows) + 1):
         messages = build_messages(case, args.prompt_style)
         prompt = render_prompt(generator, messages)
@@ -139,6 +137,14 @@ def main() -> int:
             "provider_metadata": {"model": args.model},
         }
         rows.append(row)
+        print(
+            f"[{index:03d}/{total_cases:03d}] {case.case_id} "
+            f"token_f1={row['scores']['token_f1']:.4f} "
+            f"rougeL={row['scores'].get('rougeL_f1') or 0:.4f} "
+            f"words={row['scores']['word_count']} "
+            f"latency_ms={row['latency_ms']:.2f}",
+            flush=True,
+        )
         if args.save_every > 0 and len(rows) % args.save_every == 0:
             write_progress(output_path, rows, argparse.Namespace(
                 provider="hf-transformers",
@@ -146,14 +152,11 @@ def main() -> int:
                 dataset=args.dataset,
                 prompt_style=args.prompt_style,
             ), final=False)
-            print(f"Saved progress at {len(rows)}/{len(cases)} cases -> {output_path}")
+            print(f"Saved progress at {len(rows)}/{total_cases} cases -> {output_path}", flush=True)
         if args.verbose:
             print(
-                f"[{index:03d}] {case.case_id} "
-                f"token_f1={row['scores']['token_f1']:.4f} "
-                f"rougeL={row['scores'].get('rougeL_f1') or 0:.4f} "
-                f"words={row['scores']['word_count']} "
-                f"latency_ms={row['latency_ms']:.2f}"
+                f"PRED: {prediction[:300]}",
+                flush=True,
             )
 
     compute_bertscore(rows, args.disable_bertscore, args.bertscore_model)