File size: 1,459 Bytes
f18dc10 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 |
import streamlit as st
import transformers
import torch
# 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() |