Update app.py

app.py

@@ -7,11 +7,12 @@ from transformers import (
     StoppingCriteria,
     StoppingCriteriaList,
 )
-from threading import Thread
+import threading
+
+import ctypes
 
 tokenizer = AutoTokenizer.from_pretrained("pfnet/plamo-2-1b", trust_remote_code=True)
 
-streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=False)
 pipeline = transformers.pipeline(
     "text-generation",
     model="pfnet/plamo-2-1b",
@@ -35,10 +36,56 @@ class StoppingCriteriaSub(StoppingCriteria):
 stopping_criteria = StoppingCriteriaList([StoppingCriteriaSub(stops=["\n\n"])])
 
 
-async def respond(prompt, max_tokens):
+class CancelableThread(threading.Thread):
+    def __init__(self, group=None, target=None, name=None, args=(), kwargs={}):
+        threading.Thread.__init__(self, group=group, target=target, name=name)
+        self.args = args
+        self.kwargs = kwargs
+        return
+
+    def run(self):
+        self.id = threading.get_native_id()
+        self._target(*self.args, **self.kwargs)
+
+    def get_id(self):
+        return self.id
+
+    def raise_exception(self):
+        thread_id = self.get_id()
+        resu = ctypes.pythonapi.PyThreadState_SetAsyncExc(
+            ctypes.c_long(thread_id), ctypes.py_object(SystemExit)
+        )
+        if resu > 1:
+            ctypes.pythonapi.PyThreadState_SetAsyncExc(ctypes.c_long(thread_id), 0)
+            print("Failure in raising exception")
+
+
+class ThreadManager:
+    def __init__(self, thread: CancelableThread, **kwargs):
+        self.thread = thread
+
+    def __enter__(self):
+        # スレッドを開始
+        self.thread.start()
+        return self.thread
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        # スレッドの終了を待機
+        if self.thread.is_alive():
+            print("trying to terminate thread")
+            self.thread.raise_exception()
+        self.thread.join()
+        print("Thread has been successfully joined.")
+
+
+def respond(prompt, max_tokens):
     print(prompt)
 
-    thread = Thread(
+    streamer = TextIteratorStreamer(
+        tokenizer, skip_prompt=True, skip_special_tokens=True
+    )
+
+    thread = CancelableThread(
         target=pipeline,
         kwargs=dict(
             text_inputs=prompt,
@@ -50,24 +97,21 @@ async def respond(prompt, max_tokens):
             eos_token_id=tokenizer.eos_token_id,
             stopping_criteria=stopping_criteria,
         ),
-        daemon=True,
     )
-    thread.start()
-
     response = ""
 
-    for output in streamer:
-        if not output:
-            await asyncio.sleep(0)
-            continue
-        print(output)
-        response += output
-        yield response, gr.update(interactive=False), gr.update(interactive=False),
-    yield (
-        response,
-        gr.update(interactive=True),
-        gr.update(interactive=True),
-    )
+    with ThreadManager(thread=thread):
+        for output in streamer:
+            if not output:
+                continue
+            print(output)
+            response += output
+            yield response, gr.update(interactive=False), gr.update(interactive=False),
+        yield (
+            response,
+            gr.update(interactive=True),
+            gr.update(interactive=True),
+        )
 
 
 def reset_textbox():