ICE3 / app.py
long1104's picture
Update app.py
93aa755 verified
raw
history blame
No virus
1.45 kB
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()