| # 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() | |
| import streamlit as st | |
| from sentence_transformers import SentenceTransformer | |
| from sklearn.metrics.pairwise import cosine_similarity | |
| # Load pre-trained Sentence Transformer model | |
| model = SentenceTransformer('bert-base-nli-mean-tokens') | |
| # Define your dataset here (example categories) | |
| categories = { | |
| 'sports': ['football', 'basketball', 'tennis'], | |
| 'politics': ['election', 'government', 'policy'], | |
| 'technology': ['AI', 'machine learning', 'data science'] | |
| } | |
| # Function to get relevant categories based on user query | |
| def get_relevant_categories(query): | |
| query_embedding = model.encode([query]) | |
| category_scores = {} | |
| for category, keywords in categories.items(): | |
| keyword_embeddings = model.encode(keywords) | |
| similarity_scores = cosine_similarity(query_embedding, keyword_embeddings) | |
| category_scores[category] = sum(similarity_scores)[0] | |
| relevant_categories = [category for category, score in sorted(category_scores.items(), key=lambda x: x[1], reverse=True) if score > 0] | |
| return relevant_categories | |
| # Streamlit app layout and UI | |
| def main(): | |
| st.title("Sentence Categorization App") | |
| st.write("Enter a sentence to categorize:") | |
| user_input = st.text_input('', value='', max_chars=None, key=None, type='default') | |
| if st.button('Categorize'): | |
| if user_input: | |
| relevant_categories = get_relevant_categories(user_input) | |
| st.write("Relevant Categories:") | |
| for category in relevant_categories: | |
| st.write(f"- {category}") | |
| else: | |
| st.write("Please enter a sentence for categorization.") | |
| if __name__ == "__main__": | |
| main() |