ML-project / app.py
Yanxi Lin
tryout: switch to another pipline
e76ca61
raw
history blame contribute delete
947 Bytes
import streamlit as st
from transformers import pipeline
from transformers import BartForConditionalGeneration, AutoTokenizer
# Create a pipeline from https://huggingface.co/s-nlp/bart-base-detox
pipe = pipeline("text2text-generation", model="s-nlp/bart-base-detox")
base_model_name = 'facebook/bart-base'
model_name = 's-nlp/bart-base-detox'
tokenizer = AutoTokenizer.from_pretrained(base_model_name)
model = BartForConditionalGeneration.from_pretrained(model_name)
# Streamlit app structure
st.title("Text2Text Generation App for Machine Learning Project")
st.write("Enter some cursing sentence (toxic sentence) and the model will generate the non-toxic one:")
# Input text from the user
user_input = st.text_input("Input sentence:")
# When the button is clicked, classify the input
if st.button("Generate"):
if user_input:
output = pipe(user_input)
st.write(output)
else:
st.write("Please enter some text.")