Update app.py

app.py

@@ -2,17 +2,16 @@
 
 import os
 from threading import Thread
+from queue import Queue, Empty
 from typing import Iterator
 
 import gradio as gr
 import spaces
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
 DESCRIPTION = "# Sakaltum-7B-chat"
-
-if not torch.cuda.is_available():
-    DESCRIPTION += "\n<p>Running on CPU 🥶 This demo might be slower on CPU.</p>"
+DESCRIPTION += "\n<p>現在の環境に合わせて最適化されています。</p>"
 
 MAX_MAX_NEW_TOKENS = 2048
 DEFAULT_MAX_NEW_TOKENS = 1024
@@ -22,7 +21,7 @@ model_id = "sakaltcommunity/sakaltum-7b"
 if torch.cuda.is_available():
     model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype="auto", device_map="auto")
 else:
-    model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16)
+    model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float32)
 model.eval()
 
 tokenizer = AutoTokenizer.from_pretrained(model_id)
@@ -56,26 +55,34 @@ def generate(
         gr.Warning(f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")
     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,
-        max_new_tokens=max_new_tokens,
-        do_sample=True,
-        top_p=top_p,
-        top_k=top_k,
-        temperature=temperature,
-        num_beams=1,
-        repetition_penalty=repetition_penalty,
-        pad_token_id=tokenizer.eos_token_id,
-    )
-    t = Thread(target=model.generate, kwargs=generate_kwargs)
-    t.start()
+    output_queue = Queue()
+    def inference():
+        outputs = model.generate(
+            input_ids=input_ids,
+            max_new_tokens=max_new_tokens,
+            do_sample=True,
+            top_p=top_p,
+            top_k=top_k,
+            temperature=temperature,
+            repetition_penalty=repetition_penalty,
+            pad_token_id=tokenizer.eos_token_id,
+        )
+        for token in tokenizer.decode(outputs[0], skip_special_tokens=True).split():
+            output_queue.put(token)
+        output_queue.put(None)  # 終了シグナル
+
+    Thread(target=inference).start()
 
     outputs = []
-    for text in streamer:
-        outputs.append(text)
-        yield "".join(outputs)
+    while True:
+        try:
+            token = output_queue.get(timeout=20.0)  # タイムアウト設定
+            if token is None:
+                break
+            outputs.append(token)
+            yield "".join(outputs)
+        except Empty:
+            yield "現在応答を生成中です。しばらくお待ちください。"
 
 
 demo = gr.ChatInterface(