Update app.py

app.py

@@ -1,56 +1,33 @@
-from transformers import GPTNeoForCausalLM, GPT2Tokenizer
 import gradio as gr
 
-model = GPTNeoForCausalLM.from_pretrained("EleutherAI/gpt-neo-125M")
-tokenizer = GPT2Tokenizer.from_pretrained("EleutherAI/gpt-neo-125M")
-
-prompt = """This is a discussion between a person and Hassan Kane, an entrepreneur. 
-person: What are you working on?
-Hassan: This new AI community building the future of Africa
-person: Where are you?
-Hassan: In Lagos for a week, then Paris or London.
-person: How's it going? 
-Hassan: Not bad.. Just trying to hit EV (escape velocity) with my startup
-person: """
-
-def my_split(s, seps):
-    res = [s]
-    for sep in seps:
-        s, res = res, []
-        for seq in s:
-            res += seq.split(sep)
-    return res
-
-# input = "Who are you?"
-def chat_base(input):
-  p = prompt + input
-  input_ids = tokenizer(p, return_tensors="pt").input_ids
-  gen_tokens = model.generate(input_ids, do_sample=True, temperature=0.7, max_length=150,)
-  gen_text = tokenizer.batch_decode(gen_tokens)[0]
-  # print(gen_text)
-  result = gen_text[len(p):]   
-  # print(">", result)
-  result = my_split(result, [']', '\n'])[1]
-  # print(">>", result)
-  if "Hassan: " in result:
-   result = result.split("Hassan: ")[-1]
-# print(">>>", result)
-  return result
-  
-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>"
+
+
+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
+     
 
-def chat(message):
-    history = gr.get_state() or []
-    print(history)
-    response = chat_base(message)
-    history.append((message, response))
-    gr.set_state(history)
-    html = "<div class='chatbot'>"
-    for user_msg, resp_msg in history:
-        html += f"<div class='user_msg'>{user_msg}</div>"
-        html += f"<div class='resp_msg'>{resp_msg}</div>"
-    html += "</div>"
-    return response
-
-iface = gr.Interface(chat_base, gr.inputs.Textbox(label="Ask Hassan a Question"), "text", allow_screenshot=False, allow_flagging=False)
-iface.launch()
+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)