import gradio as gr
from load_model import pipe
from utils import Translation



def reply(message, history):
    txt = Translation(message, "en")
    if txt.original == "en":
        response = pipe(message)
        return response[0]["generated_text"].split("Answer\n")[1]
    else:
        translation = txt.translatef()
        response = pipe(translation)
        t = Translation(response[0]["generated_text"].split("Answer\n")[1], txt.original)
        res = t.translatef()
        return res


demo = gr.ChatInterface(fn=reply, examples=["What is data science?", "你好你好吗？", "Ποιος ήταν ο Αρχιμήδης;", "महात्मा गांधी के बारे में बात करें", "Descrivi la miglior ricetta italiana", "Wo liegt die Mauer von Berlin?"], title="Multilingual-Phi Bot")
demo.launch()