import streamlit as st
import transformers

# Load the pre-trained language model
model_name = "bert-base-uncased"
model = transformers.pipeline("text-classification", model=model_name)

# Streamlit App
def main():
    st.title("Sentence Category Classifier")
    
    # Input search sentence
    search_query = st.text_input("Enter a sentence:")
    
    result = ""
    
    # Process the search sentence when the user clicks the Search button
    if st.button("Search"):
        if search_query:
            # Classify the sentence using the pre-trained model
            categories = classify_sentence(search_query)
            
            # Display the categories as output
            if categories:
                result = f"The sentence belongs to the following categories:\n\n"
                for category in categories:
                    result += f"• {category}\n"
            else:
                result = "No categories found for the sentence."
    
    # Display the result
    st.text(result)

# Function to classify the sentence using the pre-trained language model
@st.cache(allow_output_mutation=True)
def classify_sentence(query):
    # Classify the sentence using the pre-trained model
    categories = model(query)
    
    # Extract the category labels from the model's output
    category_labels = [category['label'] for category in categories]
    
    return category_labels

if __name__ == "__main__":
    main()