import gradio as gr 
from transformers import Conversation, pipeline
from typing import Union

chatbot = pipeline("conversational", model="facebook/blenderbot-400M-distill")

def process(prompt:Union[str,None]=None,history:Union[str,None]=None):
    """" process prompt
    history has the following format 
    ```
    ****user \t hello
    ****assistant \t hi
    ****user \t how are you
    ****assistant \t good
    """
    conv = Conversation()
    if history != None : 
        for hist in history.split("****") : 
            if len(hist.strip()) == 0 : continue
            entry = hist.split("\t")
            conv.add_message({"role": entry[0].strip(), "content": entry[1].strip()})
    if prompt is not None or prompt != "":
        conv.add_message({"role": "user", "content": prompt})
    return chatbot(conv).messages[-1]["content"]


prompt = gr.Textbox(label="Prompt")
history = gr.Textbox(label="History")
out = gr.Textbox(label="Response")
examples = [
    ["",""" ****user \t hello \n ***assistant \t hi \n ****user \t how are you"""],
    ["what's 2+10",""" ****user \t hi \n ***assistant \t hello \n"""]
    ]

title = "Backend For Conversational Bot"
description = """this is the backend for a conversational discord bot
this space is running on CPU , click here to duplicate it <a style="display:inline-block" href="https://huggingface.co/spaces/not-lain/backend-for-conversational-bot?duplicate=true"><img src="https://img.shields.io/badge/-Duplicate%20Space-blue?labelColor=white&style=flat&logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABAAAAAQCAYAAAAf8/9hAAAAAXNSR0IArs4c6QAAAP5JREFUOE+lk7FqAkEURY+ltunEgFXS2sZGIbXfEPdLlnxJyDdYB62sbbUKpLbVNhyYFzbrrA74YJlh9r079973psed0cvUD4A+4HoCjsA85X0Dfn/RBLBgBDxnQPfAEJgBY+A9gALA4tcbamSzS4xq4FOQAJgCDwV2CPKV8tZAJcAjMMkUe1vX+U+SMhfAJEHasQIWmXNN3abzDwHUrgcRGmYcgKe0bxrblHEB4E/pndMazNpSZGcsZdBlYJcEL9Afo75molJyM2FxmPgmgPqlWNLGfwZGG6UiyEvLzHYDmoPkDDiNm9JR9uboiONcBXrpY1qmgs21x1QwyZcpvxt9NS09PlsPAAAAAElFTkSuQmCC&logoWidth=14" alt="Duplicate Space"></a>

"""

iface = gr.Interface(fn=process,inputs=[prompt,history],outputs=[out],examples=examples,cache_examples=True,title=title,description=description)
iface.launch()