digitalWDF / app.py
bigPear's picture
Update app.py
8c4e6f3
raw
history blame
1.26 kB
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()
import time
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:
gr.Markdown('''## DigitalWDF - unofficial demo
''')
chatbot = gr.Chatbot([], elem_id="chatbot").style(height=200)
def user(user_message, history):
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()