Update app.py

app.py

@@ -1,17 +1,14 @@
 import os
-from collections.abc import Iterator
-from threading import Thread
-
 import gradio as gr
 import spaces
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
 DESCRIPTION = """\
 # EvaByte-SFT
 
-EvaByte is a byte-level language model that combines multibyte prediction with the efficient EVA attention mechanism.
-This page hosts [`EvaByte/EvaByte-SFT`](https://huggingface.co/EvaByte/EvaByte-SFT), fine-tuned via supervised instruction data to enable chat and general instruction-following capabilities.  
+EvaByte is a byte-level language model that combines multibyte prediction with the efficient EVA attention mechanism.  
+This page hosts [EvaByte/EvaByte-SFT](https://huggingface.co/EvaByte/EvaByte-SFT), fine-tuned via supervised instruction data to enable chat and general instruction-following capabilities.  
 For full details on architecture, training recipe, and benchmarks, see their blog post and the project repository:
 
 - Blog: <https://hkunlp.github.io/blog/2025/evabyte>  
@@ -25,71 +22,61 @@ MAX_INPUT_TOKEN_LENGTH = 32000
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
 tokenizer = AutoTokenizer.from_pretrained("EvaByte/EvaByte", trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained("evabyte/EvaByte-SFT", torch_dtype=torch.bfloat16, trust_remote_code=True).eval().to("cuda")
+model = AutoModelForCausalLM.from_pretrained(
+    "EvaByte/EvaByte-SFT",
+    torch_dtype=torch.bfloat16,
+    trust_remote_code=True,
+).eval().to(device)
+
 
 @spaces.GPU(duration=90)
 def generate(
     message: str,
     chat_history: list[dict],
-    max_new_tokens: int = 1024,
+    max_new_tokens: int = DEFAULT_MAX_NEW_TOKENS,
     temperature: float = 0.6,
     top_p: float = 0.9,
-) -> Iterator[str]:
+) -> str:                                  
+
     conversation = [*chat_history, {"role": "user", "content": message}]
+    input_ids = tokenizer.apply_chat_template(
+        conversation,
+        add_generation_prompt=True,
+        return_tensors="pt"
+    )
 
-    input_ids = tokenizer.apply_chat_template(conversation, add_generation_prompt=True, return_tensors="pt")
     if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
         input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
-        gr.Warning(f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")
+        gr.Warning(
+            f"Trimmed input to the last {MAX_INPUT_TOKEN_LENGTH} tokens　because it exceeded the limit."
+        )
+
     input_ids = input_ids.to(model.device)
 
-    streamer = TextIteratorStreamer(tokenizer, timeout=20.0, skip_prompt=True, skip_special_tokens=True)
-    generate_kwargs = dict(
-        {"input_ids": input_ids},
-        #streamer=streamer,
+    output_ids = model.multi_byte_generate(  
+        input_ids=input_ids,
         max_new_tokens=max_new_tokens,
         do_sample=True,
         top_p=top_p,
         temperature=temperature,
     )
-    t = Thread(target=model.multi_byte_generate, kwargs=generate_kwargs)
-    t.start()
 
-    outputs = []
-    for text in streamer:
-        outputs.append(text)
-        yield "".join(outputs)
+    generated_segment = output_ids[0][input_ids.shape[1]:]
+    return tokenizer.decode(generated_segment, skip_special_tokens=False, clean_up_tokenization_spaces=False)
 
 
 demo = gr.ChatInterface(
     fn=generate,
     additional_inputs=[
-        gr.Slider(
-            label="Max new tokens",
-            minimum=1,
-            maximum=MAX_MAX_NEW_TOKENS,
-            step=1,
-            value=DEFAULT_MAX_NEW_TOKENS,
-        ),
-        gr.Slider(
-            label="Temperature",
-            minimum=0.1,
-            maximum=4.0,
-            step=0.1,
-            value=0.6,
-        ),
-        gr.Slider(
-            label="Top-p (nucleus sampling)",
-            minimum=0.05,
-            maximum=1.0,
-            step=0.05,
-            value=0.9,
-        ),
-    ],
-    stop_btn=None,
-    examples=[
-        ["Write me an English pangram."],
+        gr.Slider("Max new tokens", minimum=1, maximum=MAX_MAX_NEW_TOKENS,
+                  step=1, value=DEFAULT_MAX_NEW_TOKENS),
+        gr.Slider("Temperature", minimum=0.1, maximum=4.0,
+                  step=0.1, value=0.6),
+        gr.Slider("Top-p (nucleus sampling)", minimum=0.05,
+                  maximum=1.0, step=0.05, value=0.9),
     ],
+    stop_btn=None,           
+    examples=[["Write me an English pangram."]],
     cache_examples=False,
     type="messages",
     description=DESCRIPTION,