import streamlit as st
from transformers import pipeline


@st.cache(allow_output_mutation=True, show_spinner=False)
def load_pipe():
    pipe = pipeline("text2text-generation", model="maximedb/reviews-generator")
    return pipe


st.title("Reviews Generator")
st.subheader("Pick a rating")

st.sidebar.header("Settings")
st.sidebar.subheader("Edit generate settings")
max_length = st.sidebar.slider("Max Length", min_value=10, max_value=64, value=32)
temperature = st.sidebar.slider("Temperature", value=1.0, min_value=0.0, max_value=1.0, step=0.05)
top_k = st.sidebar.slider("Top-k", min_value=10, max_value=500, value=50)
top_p = st.sidebar.slider("Top-p", min_value=0.0, max_value=1.0, step=0.05, value=1.0)

# Loading model
with st.spinner('Loading model...'):
    pipe = load_pipe()

rating = st.slider("Rating", min_value=1, max_value=5, value=3)

if st.button("Generate"):
    with st.spinner('Generating...'):
        generated = pipe(str(rating), do_sample=True, max_length=max_length, temperature=temperature, top_k=top_k, top_p=top_p)[0]["generated_text"]
    st.success(generated)