# streamlit for gpt2 model web app
import streamlit as st
import tensorflow as tf
from transformers import TFGPT2LMHeadModel, GPT2Tokenizer

tokenizer = GPT2Tokenizer.from_pretrained("ashiqabdulkhader/GPT2-Poet")
model = TFGPT2LMHeadModel.from_pretrained("ashiqabdulkhader/GPT2-Poet")

st.title("GPT2 Poet")
st.write("This is a web app for GPT2 Poet model. You can generate poems using this web app.")

prompt = st.text_input("Enter a prompt for the poem", "The quick brown fox")
length = st.slider("Length of the poem", min_value=100,
                   max_value=1000, value=100)
temperature = st.slider("Temperature", min_value=0.0, max_value=1.0, value=1.0)
top_k = st.slider("Top K", min_value=0, max_value=10, value=0)
top_p = st.slider("Top P", min_value=0.0, max_value=1.0, value=0.9)

input_ids = tokenizer.encode(prompt, return_tensors='tf')
sample_outputs = model.generate(
    input_ids,
    do_sample=True,
    max_length=length,
    top_k=top_k,
    top_p=top_p,
    temperature=temperature,
    num_return_sequences=3
)

st.write("Output:", tokenizer.decode(
    sample_outputs[0], skip_special_tokens=True))