Update app.py

app.py

@@ -1,33 +1,44 @@
-import gradio as gr
-
-title = "ELECTRA"
-
-description = "Gradio Demo for ELECTRA. To use it, simply add your text, or click one of the examples to load them. Read more at the links below."
-
-article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2003.10555' target='_blank'>ELECTRA: Pre-training Text Encoders as Discriminators Rather Than Generators</a></p>"
-
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+
+tokenizer = AutoTokenizer.from_pretrained("Evelyn18/distilbert-base-uncased-finetuned-squad")
+model = AutoModelForCausalLM.from_pretrained("Evelyn18/distilbert-base-uncased-finetuned-squad")
+
+def predict(input, history=[]):
+    # 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=1000, 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|>")
+    response.remove("")
+    
+    # write some HTML
+    html = "<div class='chatbot'>"
+    for m, msg in enumerate(response):
+        cls = "user" if m%2 == 0 else "bot"
+        html += "<div class='msg {}'> {}</div>".format(cls, msg)
+    html += "</div>"
+    
+    return html, history
 
-examples = [
-    ["My name is Sarah and I live in London","electra_large_discriminator_squad2_512","Where do I live?"]
-]
-
-io1 = gr.Interface.load("huggingface/ahotrod/electra_large_discriminator_squad2_512")
-
-io2 = gr.Interface.load("huggingface/deepset/electra-base-squad2")
-
-def inference(context, model,question):
-    if model == "electra_large_discriminator_squad2_512":
-        outlabel = io1(context,question)
-    else:
-        outlabel = io2(context,question)
-    return outlabel
-     
+import gradio as gr
 
-gr.Interface(
-    inference, 
-    [gr.inputs.Textbox(label="Context",lines=10),gr.inputs.Dropdown(choices=["electra_large_discriminator_squad2_512","electra-base-squad2"], type="value", default="electra_large_discriminator_squad2_512", label="model"),gr.inputs.Textbox(label="Question Answering")], 
-    [gr.outputs.Textbox(label="Output")],
-    examples=examples,
-    article=article,
-    title=title,
-    description=description).launch(enable_queue=True, cache_examples=True)
+css = """
+.chatbox {display:flex;flex-direction:column}
+.msg {padding:4px;margin-bottom:4px;border-radius:4px;width:80%}
+.msg.user {background-color:cornflowerblue;color:white}
+.msg.bot {background-color:lightgray;align-self:self-end}
+.footer {display:none !important}
+"""
+
+gr.Interface(fn=predict,
+             theme="default",
+             inputs=[gr.inputs.Textbox(placeholder="How are you?"), "state"],
+             outputs=["html", "state"],
+             css=css).launch()