Update app.py

app.py

@@ -1,65 +1,39 @@
-from transformers import AutoModel, AutoTokenizer
-import gradio as gr
-
-tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True)
-model = AutoModel.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True)
-#.half().cuda()
+import sys
+sys.path.append("src")
+from src import load_pretrained, ModelArguments
+
+# if __name__ == "__main__":
+model_args = ModelArguments(checkpoint_dir="temp")
+model, tokenizer = load_pretrained(model_args)
+model = model.half().cuda()
 model = model.eval()
 
-# def add_text(history, text):
-#     history = history + [(text, None)]
-#     return history, ""
-
-# def add_file(history, file):
-#     history = history + [((file.name,), None)]
-#     return history
-
-# def bot(history):
-#     response = "**That's cool!**"
-#     history[-1][1] = response
-#     return history
+import time
 
-# def predict(input, history=None):
-#     if history is None:
-#         history = []
-#     response, history = model.chat(tokenizer, input, history)
-#     return response, history
+def predict(input, history=None):
+    if history is None:
+        history = []
+    response, _ = model.chat(tokenizer, input, history)
+    history[-1][1] = ""
+    for character in response:
+        history[-1][1] += character
+        time.sleep(0.05)
+        yield history
 
 with gr.Blocks() as demo:
-    chatbot = gr.Chatbot()
-    msg = gr.Textbox()
-    clear = gr.Button("Clear")
-
-    def predict(input, history=None):
-        if history is None:
-            history = []
-        response, history = model.chat(tokenizer, input, history)
-        return response, history
-
+    gr.Markdown('''## DigitalWDF - unofficial demo
+    ''')
+    chatbot = gr.Chatbot([], elem_id="chatbot").style(height=200)
+    
+    
     def user(user_message, history):
-        return "", history + [[user_message, None]]
-
-    def bot(msg, history):
-        bot_message, _ = predict(msg, history)
-        history[-1][1] = ""
-        for character in bot_message:
-            history[-1][1] += character
-            time.sleep(0.05)
-            yield history
-
-    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
-        bot, [msg, chatbot], [msg, chatbot]
-    )
-    clear.click(lambda: None, None, chatbot, queue=False)
-
-demo.queue()
-demo.launch()
-
-    # txt.submit(predict, [txt, state], [chatbot, state])
-    # button.click(predict, [txt, state], [chatbot, state])
-
-    # btn.upload(add_file, [chatbot, btn], [chatbot]).then(
-    #     bot, chatbot, chatbot
-    # )
-
-demo.launch()
+        return history + [[user_message, None]]
+    
+    with gr.Row():
+        with gr.Column(scale=4):
+            txt = gr.Textbox(show_label=False, placeholder="Enter text and press enter").style(container=False)
+        with gr.Column(scale=1):
+            button = gr.Button("Generate")
+    txt.submit(user, [txt, chatbot], chatbot, queue=False).then(predict, [txt, chatbot], chatbot)
+    button.click(user, [txt, chatbot], chatbot, queue=False).then(predict, [txt, chatbot], chatbot)
+demo.queue().launch()