Use interface by @theodotus

app.py

@@ -1,34 +1,8 @@
 import ctranslate2
 from transformers import AutoTokenizer
 
-import threading
 import gradio as gr
 
-from typing import Optional
-from queue import Queue
-
-
-
-
-class TokenIteratorStreamer:
-    def __init__(self, end_token_id: int, timeout: Optional[float] = None):
-        self.end_token_id = end_token_id
-        self.queue = Queue()
-        self.timeout = timeout
-
-    def put(self, token: int):
-        self.queue.put(token, timeout=self.timeout)
-
-    def __iter__(self):
-        return self
-
-    def __next__(self):
-        token = self.queue.get(timeout=self.timeout)
-        if token == self.end_token_id:
-            raise StopIteration()
-        else:
-            return token
-
 
 
 def generate_prompt(history):
@@ -39,13 +13,8 @@ def generate_prompt(history):
     tokens = tokenizer.convert_ids_to_tokens(tokenizer.encode(prompt))
     return tokens
 
-def generate(streamer, history):
-    def stepResultCallback(result):
-        streamer.put(result.token_id)
-        if result.is_last and (result.token_id != end_token_id):
-            streamer.put(end_token_id)
-        print(f"step={result.step}, batch_id={result.batch_id}, token={result.token}")
-    
+def generate(question):
+    history = [[question, ""]]
     tokens = generate_prompt(history)
 
     results = translator.translate_batch(
@@ -53,43 +22,19 @@ def generate(streamer, history):
         beam_size=1,
         max_decoding_length = 256,
         repetition_penalty = 1.8,
-        callback = stepResultCallback
     )
-    return results
+    answer = tokenizer.convert_tokens_to_string(results[0].hypotheses[0])
+    return answer
 
 
 
-translator = ctranslate2.Translator("model", intra_threads=2)
+translator = ctranslate2.Translator("model")
 tokenizer = AutoTokenizer.from_pretrained("DKYoon/mt5-xl-lm-adapt")
 end_token = "</s>"
 end_token_id = tokenizer.encode(end_token)[0]
 
 
-with gr.Blocks() as demo:
-    chatbot = gr.Chatbot()
-    msg = gr.Textbox()
-    clear = gr.Button("Clear")
-
-    def user(user_message, history):
-        return "", history + [[user_message, ""]]
-
-    def bot(history):
-        bot_message_tokens = []
-        streamer = TokenIteratorStreamer(end_token_id = end_token_id)
-        generation_thread = threading.Thread(target=generate, args=(streamer, history))
-        generation_thread.start()
-        
-        for token in streamer:
-            bot_message_tokens.append(token)
-            history[-1][1] = tokenizer.decode(bot_message_tokens)
-            yield history
-        generation_thread.join()
-
-    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
-        bot, chatbot, chatbot
-    )
-    clear.click(lambda: None, None, chatbot, queue=False)
-    
-demo.queue()
+demo = gr.Interface(fn=generate, description="Space by @theodotus. Source: https://huggingface.co/spaces/theodotus/pythia-uk", inputs="text", outputs="text")
+demo.queue(1)
 if __name__ == "__main__":
     demo.launch()