import streamlit as st
from transformers import AutoTokenizer, T5ForConditionalGeneration
  
model_name = "allenai/t5-small-squad2-question-generation"
tokenizer = AutoTokenizer.from_pretrained(model_name)
@st.cache
def load_model(model_name):
    model = T5ForConditionalGeneration.from_pretrained(model_name)
    return model

model = load_model(model_name)


def run_model(input_string, **generator_args):
    input_ids = tokenizer.encode(input_string, return_tensors="pt")
    res = model.generate(input_ids, **generator_args)
    output = tokenizer.batch_decode(res, skip_special_tokens=True)
    print(output)
    return output


default_value = "Nicejob has increased our revenue 80% since signing up"

#prompts
st.title("Question Generation")
st.write("Placeholder for some other texts, like instructions...")

sent = st.text_area("Text", default_value, height = 100)
max_length = st.sidebar.slider("Max Length", min_value = 10, max_value=150,value=80,step=5)
temperature = st.sidebar.slider("Temperature", value = 1.0, min_value = 0.0, max_value=1.0, step=0.05)
num_return_sequences = st.sidebar.slider("Num Return Sequences", min_value = 1, max_value=10, value = 2)
num_beams = st.sidebar.slider("Num Beams", min_value = 1, max_value=10, value = 4)
top_k = st.sidebar.slider("Top-k", min_value = 0, max_value=100, value = 90)
top_p = st.sidebar.slider("Top-p", min_value = 0.0, max_value=1.0, step = 0.05, value = 0.9)



output_sequences = run_model(sent, max_length=max_length,num_return_sequences=num_return_sequences,
                        num_beams=num_beams,
                        temperature=temperature, top_k=top_k, top_p=top_p)

st.write(output_sequences)