File size: 374 Bytes
677039b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 |
import streamlit as st
from transformers import pipeline
task = st.selectbox("Choose a task", ["Text generation", "Summarization"])
if task == "Text generation":
model = pipeline("text-generation")
else:
model = pipeline("summarization")
prompt = st.text_area("Enter a prompt")
if prompt:
output = model(prompt)[0]["generated_text"]
st.write(output)
|