Update app.py

app.py

@@ -1,6 +1,7 @@
 import os
 from threading import Thread
 from typing import Iterator
+import queue
 
 import gradio as gr
 import spaces
@@ -29,7 +30,6 @@ model = AutoModelForCausalLM.from_pretrained(
 model.config.sliding_window = 4096
 model.eval()
 
-
 @spaces.GPU
 def generate(
     message: str,
@@ -46,11 +46,11 @@ def generate(
     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.")
-    input_ids = input_ids.to(model.device)
+    input_ids = input_ids.to(device)
 
     streamer = TextIteratorStreamer(tokenizer, timeout=20.0, skip_prompt=True, skip_special_tokens=True)
     generate_kwargs = dict(
-        {"input_ids": input_ids},
+        input_ids=input_ids,
         streamer=streamer,
         max_new_tokens=max_new_tokens,
         do_sample=True,
@@ -64,11 +64,15 @@ def generate(
     t.start()
 
     outputs = []
-    for text in streamer:
-        outputs.append(text)
+    try:
+        for text in streamer:
+            outputs.append(text)
+            yield "".join(outputs)
+    except queue.Empty:
+        # キューが空になった場合の処理
+        gr.Warning("生成プロセスがタイムアウトしました。")
         yield "".join(outputs)
 
-
 demo = gr.ChatInterface(
     fn=generate,
     type="messages",
@@ -122,6 +126,5 @@ demo = gr.ChatInterface(
     cache_examples=False,
 )
 
-
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()