import os
os.system('pip install gradio==2.3.5b0')

import gradio as gr

from transformers import pipeline

import wikipedia


talk = 0
su = ''
nlp_qa = pipeline('question-answering')

def chat(message):
    history = gr.get_state() or []


    global talk
    global su

    if talk == 0:
      response = 'tell me a subject you wanted to talk about'
      talk = 1
    
    elif talk == 1:
      results = wikipedia.search(message)
      summary = wikipedia.summary(results)
      su = summary
      response = 'ask me a question about this topic'
      talk = 2
    
    else:
      a = nlp_qa(context=su, question=message)
      response = list(a.values())[3]
      talk = 0



    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 html

iface = gr.Interface(chat, "text", "html", css="""
    .chatbox {display:flex;flex-direction:column}
    .user_msg, .resp_msg {padding:4px;margin-bottom:4px;border-radius:4px;width:80%}
    .user_msg {background-color:cornflowerblue;color:white;align-self:start}
    .resp_msg {background-color:lightgray;align-self:self-end}
""", allow_screenshot=False, allow_flagging=False)
if __name__ == "__main__":
    iface.launch(debug=True)