import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "dwojcik/gpt2-large-fine-tuned-context-256"
# model_name = "gpt2-large"
model = AutoModelForCausalLM.from_pretrained(model_name)
model.generation_config.temperature = 2.0
tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side="right")
tokenizer.pad_token = tokenizer.eos_token

def generate_response(user_message):
    inputs = tokenizer.encode(user_message, return_tensors='pt')
    outputs = model.generate(inputs, max_length=150, num_return_sequences=1)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response

def user(user_message, history):
    return gr.update(value="", interactive=False), history + [[user_message, None]]

def bot(history):
    user_message = history[-1][0]
    bot_message = generate_response(user_message)
    history[-1][1] = bot_message
    return history

with gr.Blocks() as demo:
    gr.Markdown("""
    # GPT-PTZE
    This chatbot utilizes a fine-tuned GPT-2 large model from OpenAI to generate contextually relevant responses based on user input. It was trained on large corpus of data from Przegląd Elektrotechniczny.""")
    chatbot = gr.Chatbot()
    msg = gr.Textbox("The most interesting topic in electromagnetic research is", label="Your input")
    clear = gr.Button("Clear")

    response = msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
        bot, chatbot, chatbot
    )
    response.then(lambda: gr.update(interactive=True), None, [msg], queue=False)
    clear.click(lambda: None, None, chatbot, queue=False)

demo.queue()
demo.launch(server_name="0.0.0.0")