import gradio as gr
from huggingface_hub import InferenceClient
from deep_translator import GoogleTranslator

"""
For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
"""
client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
translator_vi2en = GoogleTranslator(source='vi', target='en')
translator_en2vi = GoogleTranslator(source='en', target='vi')


def respond(
    message,
    history: list[tuple[str, str]],
    system_message,
    max_tokens,
    temperature,
    top_p,
):

    message_en = translator_vi2en.translate(message)
    
    messages = [{"role": "system", "content": system_message}]

    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})

    messages.append({"role": "user", "content": message_en})

    respnse =  client.text_generation(
        messages,
        temperature=temperature,
        top_p=top_p,
    )

    response_vi = translator_en2vi.translate(response)
        
    return response_vi


"""
For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
"""
demo = gr.ChatInterface(
    respond
)


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