import os
import gradio as gr
from gradio_client import Client

MY_HF_TOKEN_KEY = os.environ["MY_HF_TOKEN_KEY"]

client = Client("Ghana-NLP/Khaya-chat-backend",hf_token=MY_HF_TOKEN_KEY)

def generate(prompt,history,language):
    print(history)
    output = client.predict(prompt,language)
    print(output)
    history.append([prompt,output])
    return  output

with gr.Blocks(fill_height=True) as demo:
    title = gr.Markdown(
    """
    # African Language Chatbot (Khaya)
    SELECT YOUR LANGUAGE BELOW!
    Based on the <a href="https://translation.ghananlp.org/">Khaya AI</a> and <a href="https://lesan.ai/">Lesan AI</a>. Lesan is used for <b>Amharic</b> and <b>Tigrinya</b>, Khaya is used for the rest.
    """)
    language_selector = gr.Dropdown(["Amharic","Twi","Dagbani","Ewe","Ga","Gurene","Fante","Hausa", "Kikuyu", "Kimeru", "Luo","Shona","Swahili","Tigrinya","Yoruba"],value="Twi",label="Choose Language! (default is Twi)", info="Language:")
    main_block = gr.ChatInterface(fn=generate, additional_inputs = language_selector,
        description = 'Select an African language above and type messages below in selected language')

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