| import streamlit as st |
| import torch |
| import pandas as pd |
| from transformers import AutoTokenizer, AutoModelForSequenceClassification |
|
|
| MODEL_ID = "dejanseo/ecommerce-query-volume-classifier" |
|
|
| LABELS = ["very_low", "low", "medium", "high", "very_high"] |
|
|
| DISPLAY_LABELS = { |
| "very_high": "Very High", |
| "high": "High", |
| "medium": "Medium", |
| "low": "Low", |
| "very_low": "Very Low", |
| } |
|
|
| VOLUME_RANGES = { |
| "very_high": "10,000+", |
| "high": "1,000–9,999", |
| "medium": "100–999", |
| "low": "10–99", |
| "very_low": "<10", |
| } |
|
|
| CLASS_COLORS = { |
| "Very High": "background-color: rgba(40, 167, 69, 0.1); color: #28a745;", |
| "High": "background-color: rgba(113, 194, 133, 0.1); color: #71c285;", |
| "Medium": "background-color: rgba(255, 193, 7, 0.1); color: #c79100;", |
| "Low": "background-color: rgba(253, 126, 20, 0.1); color: #fd7e14;", |
| "Very Low": "background-color: rgba(220, 53, 69, 0.1); color: #dc3545;", |
| } |
|
|
|
|
| @st.cache_resource |
| def load_model(): |
| tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) |
| model = AutoModelForSequenceClassification.from_pretrained(MODEL_ID) |
| model.eval() |
| return tokenizer, model |
|
|
|
|
| def predict(queries, tokenizer, model): |
| inputs = tokenizer( |
| queries, |
| return_tensors="pt", |
| padding=True, |
| truncation=True, |
| max_length=32, |
| ) |
| with torch.no_grad(): |
| logits = model(**inputs).logits |
| probs = torch.softmax(logits, dim=-1) |
| preds = torch.argmax(probs, dim=-1) |
| results = [] |
| for i, query in enumerate(queries): |
| label = LABELS[preds[i].item()] |
| conf = probs[i][preds[i].item()].item() |
| dist = {LABELS[j]: probs[i][j].item() for j in range(len(LABELS))} |
| results.append({"query": query, "label": label, "confidence": conf, "distribution": dist}) |
| return results |
|
|
|
|
| def style_class_cells(val): |
| return CLASS_COLORS.get(val, "") |
|
|
|
|
| st.set_page_config(page_title="Query Volume Classifier", page_icon="🔍", layout="wide") |
|
|
| st.markdown(""" |
| <style> |
| @import url('https://fonts.googleapis.com/css2?family=Montserrat:wght@400;500;600;700&display=swap'); |
| html, body, [class*="css"], .stMarkdown, .stTextInput, .stTextArea, |
| .stMetric, .stDataFrame, h1, h2, h3, h4, h5, h6, p, span, div, label { |
| font-family: 'Montserrat', sans-serif !important; |
| font-weight: 500; |
| } |
| </style> |
| """, unsafe_allow_html=True) |
|
|
| st.logo( |
| "https://dejan.ai/wp-content/uploads/2024/02/dejan-300x103.png", |
| link="https://dejan.ai/", |
| size="large", |
| ) |
|
|
| st.subheader("eCommerce Query Volume Classifier") |
| st.caption("Predicts search volume class for product queries.") |
|
|
| tokenizer, model = load_model() |
|
|
| tab_single, tab_batch = st.tabs(["Single Query", "Batch"]) |
|
|
| with tab_single: |
| query = st.text_input("Enter a product search query", placeholder="e.g. wireless mouse, airpods, replacement gasket for instant pot") |
| if st.button("Classify", key="single") and query.strip(): |
| result = predict([query.strip()], tokenizer, model)[0] |
| label = result["label"] |
| display_label = DISPLAY_LABELS[label] |
| volume_range = VOLUME_RANGES[label] |
|
|
| col1, col2, col3 = st.columns(3) |
| col1.metric("Predicted Class", display_label) |
| col2.metric("Estimated Sessions", volume_range) |
| col3.metric("Confidence", f"{result['confidence']:.1%}") |
|
|
| st.markdown("**Class probabilities**") |
| for lbl in ["very_high", "high", "medium", "low", "very_low"]: |
| prob = result["distribution"][lbl] |
| st.progress(prob, text=f"{DISPLAY_LABELS[lbl]} — {prob:.1%}") |
|
|
| with tab_batch: |
| st.markdown("Enter one query per line:") |
| text = st.text_area("Queries", height=200, placeholder="laptop\nairpods\norganic flurb capsules\nreplacement gasket for instant pot duo 8 quart") |
| if st.button("Classify", key="batch") and text.strip(): |
| queries = [q.strip() for q in text.strip().splitlines() if q.strip()] |
| if len(queries) > 500: |
| st.warning("Please enter 500 queries or fewer.") |
| else: |
| results = predict(queries, tokenizer, model) |
| df = pd.DataFrame([ |
| { |
| "Query": r["query"], |
| "Class": DISPLAY_LABELS[r["label"]], |
| "Sessions": VOLUME_RANGES[r["label"]], |
| "Confidence": r["confidence"], |
| } |
| for r in results |
| ]) |
|
|
| styled = df.style.map(style_class_cells, subset=["Class"]) |
|
|
| st.dataframe( |
| styled, |
| column_config={ |
| "Confidence": st.column_config.ProgressColumn( |
| "Confidence", |
| format="%.1f%%", |
| min_value=0, |
| max_value=1, |
| ), |
| }, |
| use_container_width=True, |
| hide_index=True, |
| ) |
