Added copy of pythia-uk

app.py

@@ -0,0 +1,95 @@
+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):
+    prompt = ""
+    for chain in history[:-1]:
+        prompt += f"<human>: {chain[0]}\n<bot>: {chain[1]}\n"
+    prompt += f"<human>: {history[-1][0]}\n<bot>:"
+    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}")
+    
+    tokens = generate_prompt(history)
+
+    results = translator.translate_batch(
+        [tokens],
+        beam_size=1,
+        max_decoding_length = 256,
+        repetition_penalty = 1.8,
+        callback = stepResultCallback
+    )
+    return results
+
+
+
+translator = ctranslate2.Translator("model", intra_threads=2)
+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()
+if __name__ == "__main__":
+    demo.launch()