|
import streamlit as st |
|
import transformers |
|
import torch |
|
|
|
|
|
model_name = "bert-base-uncased" |
|
model = transformers.pipeline("text-classification", model=model_name) |
|
|
|
|
|
def main(): |
|
st.title("Sentence Category Classifier") |
|
|
|
|
|
search_query = st.text_input("Enter a sentence:") |
|
|
|
result = "" |
|
|
|
|
|
if st.button("Search"): |
|
if search_query: |
|
|
|
categories = classify_sentence(search_query) |
|
|
|
|
|
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." |
|
|
|
|
|
st.text(result) |
|
|
|
|
|
@st.cache(allow_output_mutation=True) |
|
def classify_sentence(query): |
|
|
|
categories = model(query) |
|
|
|
|
|
category_labels = [category['label'] for category in categories] |
|
|
|
return category_labels |
|
|
|
if __name__ == "__main__": |
|
main() |