Update app.py

app.py

@@ -762,7 +762,7 @@ def build_inputs(tokenizer, model_id: str, prompt: str):
 def is_diffusion_model(model_id: str) -> bool:
     return "metadiffusion" in model_id.lower()
 
-def generate_for_model(model_id: str, prompt: str) -> str:
+def generate_for_model(model_id: str, prompt: str, max_new_tokens: int = 0) -> str:
     ensure_models_loaded()
     if model_id not in models or model_id not in tokenizers:
         short = MODEL_DISPLAY.get(model_id, model_id)
@@ -773,6 +773,8 @@ def generate_for_model(model_id: str, prompt: str) -> str:
     model = models[model_id]
     cfg = GEN_DEFAULTS.get(model_id, {})
     max_new = cfg.get("max_new_tokens", 128)
+    if max_new_tokens and int(max_new_tokens) > 0:
+        max_new = max(16, min(int(max_new_tokens), 512))
     try:
         if is_diffusion_model(model_id):
             return generate_diffusion(model, tokenizer, prompt, cfg)  # type: ignore
@@ -894,6 +896,14 @@ def create_demo() -> gr.Blocks:
                     label="Your prompt",
                     placeholder="Ask anything... e.g. 'Explain quantum computing in simple terms' or 'Write a haiku about rain'",
                     lines=3,
+                    max_length=256,
+                )
+                max_new = gr.Slider(
+                    minimum=64,
+                    maximum=512,
+                    step=64,
+                    value=128,
+                    label="Response length (tokens)",
                 )
                 with gr.Row():
                     submit_btn = gr.Button("⚔️ Battle", variant="primary", scale=1)
@@ -928,6 +938,7 @@ def create_demo() -> gr.Blocks:
 
                 return {
                     "prompt": prompt,
+                    "max_new": max_new,
                     "submit_btn": submit_btn,
                     "clear_btn": clear_btn,
                     "response_a": response_a,
@@ -1012,7 +1023,7 @@ def create_demo() -> gr.Blocks:
                 return None
 
         @spaces.GPU(duration=120)
-        def battle_gpu(spec: ArenaSpec, user_prompt: str, last_pair_val):
+        def battle_gpu(spec: ArenaSpec, user_prompt: str, last_pair_val, max_new_tokens: int = 0):
             reason = moderate_prompt_impl(user_prompt)
             if reason is not None:
                 return reason, None, None, None, None
@@ -1020,11 +1031,11 @@ def create_demo() -> gr.Blocks:
             if random.random() < 0.5:
                 a, b = b, a
             ensure_models_loaded()
-            resp_a = generate_for_model(a, user_prompt)
-            resp_b = generate_for_model(b, user_prompt)
+            resp_a = generate_for_model(a, user_prompt, max_new_tokens)
+            resp_b = generate_for_model(b, user_prompt, max_new_tokens)
             return None, resp_a, resp_b, a, b
 
-        def on_submit(spec: ArenaSpec, user_prompt: str, last_pair_val):
+        def on_submit(spec: ArenaSpec, user_prompt: str, last_pair_val, max_new_tokens: int = 128):
             user_prompt = (user_prompt or "").strip()
             if not user_prompt:
                 return (
@@ -1041,7 +1052,9 @@ def create_demo() -> gr.Blocks:
                     "", "", False, user_prompt, last_pair_val,
                     leaderboard_dataframe(load_elo(spec), spec)
                 )
-            reason, resp_a, resp_b, a, b = battle_gpu(spec, user_prompt, last_pair_val)
+            if len(user_prompt) > 256:
+                user_prompt = user_prompt[:256]
+            reason, resp_a, resp_b, a, b = battle_gpu(spec, user_prompt, last_pair_val, int(max_new_tokens or 128))
             if reason is not None:
                 flag_text = f"This prompt was flagged for: {reason}"
                 return (
@@ -1196,13 +1209,13 @@ def create_demo() -> gr.Blocks:
             vote_outputs = [u["reveal_a"], u["reveal_b"], u["status"], u["vote_a"], u["vote_tie"], u["vote_both_bad"], u["vote_b"], u["new_round_btn"], u["voted_state"], b["leaderboard"]]
 
             u["submit_btn"].click(
-                fn=lambda p, lp, s=spec: on_submit(s, p, lp),
-                inputs=[u["prompt"], u["last_pair"]],
+                fn=lambda p, lp, mn, s=spec: on_submit(s, p, lp, mn),
+                inputs=[u["prompt"], u["last_pair"], u["max_new"]],
                 outputs=submit_outputs,
             )
             u["prompt"].submit(
-                fn=lambda p, lp, s=spec: on_submit(s, p, lp),
-                inputs=[u["prompt"], u["last_pair"]],
+                fn=lambda p, lp, mn, s=spec: on_submit(s, p, lp, mn),
+                inputs=[u["prompt"], u["last_pair"], u["max_new"]],
                 outputs=submit_outputs,
             )
             for btn, choice in ((u["vote_a"], "A"), (u["vote_tie"], "Tie"), (u["vote_both_bad"], "Both Bad"), (u["vote_b"], "B")):