Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| import torch | |
| from transformers import AlbertTokenizer, AlbertForSequenceClassification | |
| import plotly.graph_objects as go | |
| # URL of the logo | |
| logo_url = "https://dejan.ai/wp-content/uploads/2024/02/dejan-300x103.png" | |
| # Display the logo at the top using st.logo | |
| st.logo(logo_url, link="https://dejan.ai") | |
| # Streamlit app title and description | |
| st.title("Search Query Form Classifier") | |
| st.write( | |
| "Ambiguous search queries are candidates for query expansion. Our model identifies such queries with an 80 percent accuracy and is deployed in a batch processing pipeline directly connected with Google Search Console API. In this demo you can test the model capability by testing individual queries." | |
| ) | |
| st.write("Enter a query to check if it's well-formed:") | |
| # Load the model and tokenizer from the Hugging Face Model Hub | |
| model_name = 'dejanseo/Query-Quality-Classifier' | |
| tokenizer = AlbertTokenizer.from_pretrained(model_name) | |
| model = AlbertForSequenceClassification.from_pretrained(model_name) | |
| # Set the model to evaluation mode | |
| model.eval() | |
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
| model.to(device) | |
| # Create tabs for single and bulk queries | |
| tab1, tab2 = st.tabs(["Single Query", "Bulk Query"]) | |
| with tab1: | |
| user_input = st.text_input("Query:", "where can I book cheap flights to london") | |
| #st.write("Developed by [Dejan AI](https://dejan.ai/blog/search-query-quality-classifier/)") | |
| def classify_query(query): | |
| # Tokenize input | |
| inputs = tokenizer.encode_plus( | |
| query, | |
| add_special_tokens=True, | |
| max_length=32, | |
| padding='max_length', | |
| truncation=True, | |
| return_attention_mask=True, | |
| return_tensors='pt' | |
| ) | |
| input_ids = inputs['input_ids'].to(device) | |
| attention_mask = inputs['attention_mask'].to(device) | |
| # Perform inference | |
| with torch.no_grad(): | |
| outputs = model(input_ids, attention_mask=attention_mask) | |
| logits = outputs.logits | |
| softmax_scores = torch.softmax(logits, dim=1).cpu().numpy()[0] | |
| confidence = softmax_scores[1] * 100 # Confidence for well-formed class | |
| return confidence | |
| # Function to determine color based on confidence | |
| def get_color(confidence): | |
| if confidence < 50: | |
| return 'rgba(255, 51, 0, 0.8)' # Red | |
| else: | |
| return 'rgba(57, 172, 57, 0.8)' # Green | |
| # Check and display classification for single query | |
| if user_input: | |
| confidence = classify_query(user_input) | |
| # Plotly grey placeholder bar with dynamic color fill | |
| fig = go.Figure() | |
| # Placeholder grey bar | |
| fig.add_trace(go.Bar( | |
| x=[100], | |
| y=['Well-formedness Factor'], | |
| orientation='h', | |
| marker=dict( | |
| color='lightgrey' | |
| ), | |
| width=0.8 | |
| )) | |
| # Colored bar based on confidence | |
| fig.add_trace(go.Bar( | |
| x=[confidence], | |
| y=['Well-formedness Factor'], | |
| orientation='h', | |
| marker=dict( | |
| color=get_color(confidence) | |
| ), | |
| width=0.8 | |
| )) | |
| fig.update_layout( | |
| xaxis=dict(range=[0, 100], title='Well-formedness Factor'), | |
| yaxis=dict(showticklabels=False), | |
| width=600, | |
| height=250, # Increase height for better visibility | |
| title_text='Well-formedness Factor', | |
| plot_bgcolor='rgba(0,0,0,0)', | |
| showlegend=False | |
| ) | |
| st.plotly_chart(fig) | |
| if confidence >= 50: | |
| st.success(f"Query Score: {confidence:.2f}% Most likely doesn't require query expansion.") | |
| st.subheader(f":sparkles: What's next?", divider="gray") | |
| st.write("Connect with Google Search Console, Semrush, Ahrefs or any other search query source API and detect all queries which could benefit from expansion.") | |
| st.write("[Engage our team](https://dejan.ai/call/) if you'd like us to do this for you.") | |
| else: | |
| st.error(f"The query is likely not well-formed with a score of {100 - confidence:.2f}% and most likely requires query expansion.") | |
| st.subheader(f":sparkles: What's next?", divider="gray") | |
| st.write("Connect with Google Search Console, Semrush, Ahrefs or any other search query source API and detect all queries which could benefit from expansion.") | |
| st.write("[Engage our team](https://dejan.ai/call/) if you'd like us to do this for you.") | |
| with tab2: | |
| st.write("Paste multiple queries line-separated (no headers or extra data):") | |
| bulk_input = st.text_area("Bulk Queries:", height=200) | |
| if bulk_input: | |
| bulk_queries = bulk_input.splitlines() | |
| st.write("Processing queries...") | |
| # Classify each query in bulk input | |
| results = [(query, classify_query(query)) for query in bulk_queries] | |
| # Display results in a table | |
| for query, confidence in results: | |
| st.write(f"Query: {query} - Score: {confidence:.2f}%") | |
| if confidence >= 50: | |
| st.success("Well-formed") | |
| else: | |
| st.error("Not well-formed") | |
| st.subheader(f":sparkles: What's next?", divider="gray") | |
| st.write("Connect with Google Search Console, Semrush, Ahrefs or any other search query source API and detect all queries which could benefit from expansion.") | |
| st.write("[Engage our team](https://dejan.ai/call/) if you'd like us to do this for you.") |