import streamlit as st
from transformers import AutoModelWithLMHead, AutoTokenizer

# Title of the page
st.title("Text Generation with Huggingface Model")

# Tokenizer model selection
model_name = st.selectbox("Select a Huggingface model",
                         ["distilbert-base-cased",
                          "gpt2",
                          "xlm-roberta-base"])

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelWithLMHead.from_pretrained(model_name)

# Text box where the user can enter the text
text = st.text_area("Enter the text",
                    "Type your text here...")

# Generate the text
if st.button("Generate"):
    input_ids = tokenizer.encode(text, return_tensors="pt")
    output_ids = model.generate(input_ids)[0]
    generated_text = tokenizer.decode(output_ids)
    st.write(generated_text)