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.")