timofeyk/roberta-query-router-ecommerce

E-commerce Search Query Router

Model Description

This model is a classifier fine-tuned from a roberta-base to determine the optimal search strategy for e-commerce queries. It classifies a given query into one of two labels:

• lexical_search: Indicates that the query is best handled by a traditional, keyword-based search engine like Lucene using BM25. These are typically specific queries like SKUs, exact product names, or part numbers.

• vector_search: Indicates that the query is better suited for a semantic, vector-based search. These are often ambiguous, conceptual, or "long-tail" queries where user intent is more important than specific keywords (e.g., "a gift for my dad who likes fishing").

The model is intended to be used as an intelligent "query router" in a hybrid search system, dynamically weighting the results from lexical and vector search engines to improve relevance.

Intended Use & Use Case

The primary use case for this model is to power a hybrid search relevance system. The intended workflow is as follows:

1. A user enters a search query.
2. The query is sent to this classifier model.
3. The model returns probabilities for the lexical_search and vector_search classes.
4. These probabilities are used as weights to blend the relevance scores from two separate search backends (e.g., Solr and Qdrant).
5. The final, blended scores are used to rank the products shown to the user.

How to Use

Here's how to use the model with the transformers library pipeline:

from transformers import pipeline

router_pipeline = pipeline(
    "text-classification",
    model="timofeyk/roberta-query-router-ecommerce",
    return_all_scores=True
)

# Example of a conceptual query
conceptual_query = "father day gift"
# Example of a specific query
specific_query = "16x16 pillow cover"

queries = [conceptual_query, specific_query]

for q in queries:
    print(f"Predicting label for query: {q}")
    results = router_pipeline(q)
    print(results[0])
    # Expected output might look like:
    # [{'label': 'lexical_search', 'score': 0.46258628368377686}, {'label': 'vector_search', 'score': 0.5374137163162231}]

    scores = {item['label']: item['score'] for item in results[0]}
    w_vector = scores['vector_search']
    w_lexical = scores['lexical_search']

    print(f"Vector Search Weight: {w_vector:.2f}")
    print(f"Lexical Search Weight: {w_lexical:.2f}")

Training Data

This model was trained on a custom dataset of anonymized, real-world e-commerce queries. The dataset was generated using Amazon ESCI Dataset as a source. The labels were generated programmatically based on search performance, creating a signal for the model to learn from:

1. A large set of user queries and their corresponding ground-truth "ideal" product results (based on user engagement) were collected.
2. Each query was executed against both a Solr (lexical) and a Qdrant (vector) search engine.
3. The nDCG relevance score was calculated for both result sets against the ground truth.
4. The query was labeled lexical_search if Solr achieved a higher nDCG score, and vector_search otherwise.

Training Procedure

The model was fine-tuned using the Hugging Face Trainer. To account for a potential class imbalance in the training data, a custom Trainer with a weighted CrossEntropyLoss was used, preventing the model from favoring the majority class.

Training params

TrainingArguments(
    learning_rate=1e-05,
    lr_scheduler_type=SchedulerType.COSINE,
    max_grad_norm=1.0,
    num_train_epochs=3,
    optim=OptimizerNames.ADAMW_TORCH_FUSED,
    optim_args=None,
    per_device_eval_batch_size=128,
    per_device_train_batch_size=32,
    prediction_loss_only=False,
    warmup_ratio=0.05,
    weight_decay=0.01,
)

Citation

If you use this model in your work, please consider citing it:

@misc{timofeyk_roberta-query-router-ecommerce,
  author = {Timofey Klyubin},
  title = {E-commerce Search Query Router},
  year = {2025},
  publisher = {Hugging Face},
  journal = {Hugging Face repository},
  howpublished = {\url{[https://huggingface.co/timofeyk/roberta-query-router-ecommerce](https://huggingface.co/timofeyk/roberta-query-router-ecommerce)}}
}