ml-playground / app.py
Sasidhar's picture
Update app.py
ab5688d
raw
history blame
707 Bytes
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
@st.cache
def load_text_gen_model():
generator = pipeline(“text-generation”, model=”EleutherAI/gpt-neo-2.7B”)
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":
st.text_input("Enter a prompt")