from transformers import TFAutoModelForCausalLM, AutoTokenizer
import tensorflow as tf
import gradio as gr

# configuration params
TITLE = "<center><h1>Talk with an AI</h1></center>"

# Loading necessary NLP models
checkpoint = "elapt1c/ElapticAI-1a"  # tf
model_gtp2 = TFAutoModelForCausalLM.from_pretrained(checkpoint)
tokenizer_gtp2 = AutoTokenizer.from_pretrained(checkpoint)

# test-to-test : chatting function -- GPT2
def chat_with_bot(user_input, chat_history_and_input=[]):
    """Text generation using GPT2"""
    emb_user_input = tokenizer_gtp2.encode(
        user_input + tokenizer_gtp2.eos_token, return_tensors="tf"
    )
    if chat_history_and_input == []:
        bot_input_ids = emb_user_input  # first iteration
    else:
        bot_input_ids = tf.concat(
            [chat_history_and_input, emb_user_input], axis=-1
        )  # other iterations
    chat_history_and_input = model_gtp2.generate(
        bot_input_ids, max_length=50, pad_token_id=tokenizer_gtp2.eos_token_id
    ).numpy()
    bot_response = tokenizer_gtp2.decode(
        chat_history_and_input[:, bot_input_ids.shape[-1] :][0],
        skip_special_tokens=True,
    )
    # Limit history to last 500 characters
    chat_history_and_input = chat_history_and_input[:, -500:]
    return bot_response, chat_history_and_input

# gradio interface
blocks = gr.Blocks()
with blocks:
    session_state = gr.State([])
    gr.Markdown(TITLE)
    user_input = gr.Textbox(label="User Input")
    bot_response = gr.Textbox(label="Bot Response")
    user_input.change(
        chat_with_bot,
        inputs=[user_input, session_state],
        outputs=[bot_response, session_state],
    )

blocks.launch()