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)