Update app.py

app.py

@@ -63,11 +63,16 @@ SPACE_REPO_ID  = os.environ.get("SPACE_REPO_ID", "")
 # Generation / toggles
 ALLOW_WIKIPEDIA = False
 DEBUG = True
-MAX_NEW_TOKENS_GROUNDED = 384   # was 512
-MAX_NEW_TOKENS_FALLBACK = 192   # was 256
-
+MAX_NEW_TOKENS_GROUNDED = 512
+MAX_NEW_TOKENS_FALLBACK = 256
 MIN_USEFUL_CHARS = 260
 
+# Auto-continue if we hit the cap without EOS
+AUTO_CONTINUE = True
+AUTO_CONT_MAX_STEPS = 2        # continue up to 2 extra chunks
+AUTO_CONT_NEW_TOKENS = 256     # tokens per continuation step
+
+
 def dlog(tag, msg):
     if DEBUG: print(f"[{tag}] {msg}")
 
@@ -267,8 +272,35 @@ GEN_ARGS_FALLBACK = dict(
 
 def _generate(inputs, grounded: bool):
     args = GEN_ARGS_GROUNDED if grounded else GEN_ARGS_FALLBACK
+    in_len = inputs["input_ids"].shape[-1]
     with torch.inference_mode():
-        return model_lm.generate(**inputs, **args)
+        out = model_lm.generate(**inputs, **args)
+
+        if not AUTO_CONTINUE:
+            return out
+
+        steps = 0
+        while steps < AUTO_CONT_MAX_STEPS:
+            seq = out[0]
+            ended_with_eos = (seq[-1].item() == tokenizer_lm.eos_token_id)
+            hit_cap = (seq.shape[0] - in_len) >= args["max_new_tokens"]
+            if ended_with_eos or not hit_cap:
+                break
+
+            # continue generation from the current sequence
+            cont_inputs = {
+                "input_ids": seq.unsqueeze(0),
+                "attention_mask": torch.ones_like(seq).unsqueeze(0),
+            }
+            cont_inputs = {k: v.to(device) for k, v in cont_inputs.items()}
+            cont_args = dict(args)
+            cont_args["max_new_tokens"] = AUTO_CONT_NEW_TOKENS
+
+            out = model_lm.generate(**cont_inputs, **cont_args)
+            steps += 1
+
+        return out
+
 
 # ================== UTILITIES ==================
 _SANITIZE = re.compile(r"```.*?```|<\s*script[^>]*>.*?<\s*/\s*script\s*>", re.DOTALL|re.IGNORECASE)
@@ -986,45 +1018,43 @@ with gr.Blocks(theme="soft") as demo:
             submit_fb = gr.Button("Submit feedback")
             fb_status = gr.Markdown("")
 
-    # Wiring
-    enter_btn.click(
-        fn=enter_app,
-        inputs=[first_tb, last_tb, state],
-        outputs=[gate, app, state, gate_msg],
-    )
+# Wiring
+enter_btn.click(
+    fn=enter_app,
+    inputs=[first_tb, last_tb, state],
+    outputs=[gate, app, state, gate_msg],
+)
 
-    send_btn.click(
-    predict,
+send_btn.click(
+    fn=predict,
     inputs=[user_in, chat, state],
     outputs=[chat, user_in, feedback_grp, rating, comment, state],
-    concurrency_limit=3,
-    )
-
+    concurrency_limit=1,  # serialize LLM calls
+)
 
-    user_in.submit(
-    predict,
+user_in.submit(
+    fn=predict,
     inputs=[user_in, chat, state],
     outputs=[chat, user_in, feedback_grp, rating, comment, state],
-    concurrency_limit=3,
-    )
-
-    clear_btn.click(
-        lambda: ([], "", gr.update(visible=False), None, "", init_session()),
-        inputs=None,
-        outputs=[chat, user_in, feedback_grp, rating, comment, state],
-        concurrency_limit=4,
-    )
+    concurrency_limit=1,  # serialize LLM calls
+)
 
-    submit_fb.click(
-        fn=save_feedback,
-        inputs=[rating, comment, state],
-        outputs=[fb_status, feedback_grp],
-        concurrency_limit=4,
-    )
+clear_btn.click(
+    fn=lambda: ([], "", gr.update(visible=False), None, "", init_session()),
+    inputs=None,
+    outputs=[chat, user_in, feedback_grp, rating, comment, state],
+    concurrency_limit=4,
+)
 
+submit_fb.click(
+    fn=save_feedback,
+    inputs=[rating, comment, state],
+    outputs=[fb_status, feedback_grp],
+    concurrency_limit=4,
+)
 
+# Queue (true concurrency = 1 to prevent OOM/restarts)
+demo.queue(concurrency_count=1, max_size=64)
 
-demo.queue(max_size=64)
-demo.launch(max_threads=int(os.environ.get("MAX_THREADS", "32")))