import gradio as gr
from transformers import pipeline
from transformers import Conversation

pipe = pipeline("conversational", model="PygmalionAI/pygmalion-6b")
conversation = Conversation()
conversation.add_message({"role": "assistant", "content": "How can I help you?"})
memory = []

def run_conversation(message, history):
    #cache memory from user
    memory.append({"role": "user", "content": message})
    print(history)

    if len(history) > 0:
        for i in history:
            conversation.add_message({"role":"user", "content":i[0]})
            conversation.add_message({"role": "assistant", "content": i[1]})
    for i in memory:
        conversation.add_message(i)
    print(conversation)

    pipe(conversation)
    return conversation.generated_responses[-1]

demo = gr.ChatInterface(
        run_conversation,
        chatbot = gr.Chatbot(height=500),
        textbox = gr.Textbox(placeholder="Chat with me!", scale=7),
        title = "Test",
        description="Chat with me!"
        examples=["hello"]
        )

if __name__ == "__main__":
    demo.queue().launch()