Add file app.py

app.py

@@ -0,0 +1,82 @@
+import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+import datetime
+import __init__
+
+MODEL_NAME = __init__.MODEL_MICROSOFT_DIABLO_MEDIUM
+OUTPUT_MAX_LENGTH = __init__.OUTPUT_MAX_LENGTH
+
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
+
+
+def print_f(session_id, text):
+    print(f"{datetime.datetime.now()} | {session_id} | {text}")
+
+
+def predict(input, history, request: gr.Request):
+    session_id = 'UNKNOWN'
+    if request:
+        # Get session_id is client_ip + client_port
+        session_id = request.client.host + ':' + str(request.client.port)
+    # print_f(session_id, f" inp: {input}")
+
+    # Tokenize the new input sentence
+    new_user_input_ids = tokenizer.encode(
+        input + tokenizer.eos_token, return_tensors='pt')
+
+    # Append the new user input tokens to the chat history
+    bot_input_ids = torch.cat(
+        [torch.LongTensor(history), new_user_input_ids], dim=-1)
+
+    # Generate a response
+    history = model.generate(bot_input_ids, max_length=OUTPUT_MAX_LENGTH,
+                             pad_token_id=tokenizer.eos_token_id).tolist()
+
+    # Convert the tokens to text, and then split the responses into lines
+    response = tokenizer.decode(history[0]).split("<|endoftext|>")
+
+    # Convert to tuples of list
+    response = [(response[i], response[i + 1])
+                for i in range(0, len(response) - 1, 2)]
+
+    # Print new conversation
+    print_f(session_id, response[-1])
+
+    return response, history
+
+
+css = """
+    #row_bot{width: 70%; height: var(--size-96); margin: 0 auto}
+    #row_bot .block{background: var(--color-grey-100); height: 100%}
+    #row_input{width: 70%; margin: 0 auto}
+    #row_input .block{background: var(--color-grey-100)}
+
+    @media screen and (max-width: 768px) {
+        #row_bot{width: 100%; height: var(--size-96); margin: 0 auto}
+        #row_bot .block{background: var(--color-grey-100); height: 100%}
+        #row_input{width: 100%; margin: 0 auto}
+        #row_input .block{background: var(--color-grey-100)}    
+    }
+    """
+block = gr.Blocks(css=css, title="Chatbot")
+
+with block:
+    gr.Markdown(f"""
+        <p style="font-size:20px; text-align: center">{MODEL_NAME}</p>
+        """)
+    with gr.Row(elem_id='row_bot'):
+        chatbot = gr.Chatbot()
+    with gr.Row(elem_id='row_input'):
+        message = gr.Textbox(placeholder="Enter something")
+        state = gr.State([])
+
+        message.submit(predict,
+                       inputs=[message, state],
+                       outputs=[chatbot, state])
+        message.submit(lambda x: "", message, message)
+
+# Params ex: debug=True, share=True, server_name="0.0.0.0", server_port=5050
+block.launch()