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''' | |