import transformers
import streamlit as st

from transformers import AutoTokenizer, AutoModelWithLMHead
from transformers import pipeline

tokenizer = AutoTokenizer.from_pretrained("gpt2-large")

@st.cache
def load_model(model_name):
    model = AutoModelWithLMHead.from_pretrained("gpt2-large")
    return model
    
def load_text_gen_model():
    generator = pipeline("text-generation", model="gpt2-medium")
    return generator 
    
model = load_model("gpt2-large")

text_generator = load_text_gen_model()

action = st.sidebar.selectbox("Pick an Action", ["Generate an Article","Create an Image"])

if action == "Generate an Article":
    prompt = st.text_input("Enter a prompt")
    res = generator( prompt, max_length=100, temperature=0.7)
    st.write(res)