File size: 537 Bytes
677039b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
40f480d
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
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'''