Text2Text / app.py
0xrizzler's picture
updated comments for better understanding
40f480d verified
raw
history blame contribute delete
537 Bytes
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)
'''This sets up a simple interface with a dropdown to select either text generation or summarization,
a text area for input, and displays the generated output'''