Model Card for Model ID

Provides a relevance score for a given search query and a product (with title and description).

Can be used as retriever/ranker for e-commerce search.

Development details

Model is trained with a novel adversarial Generator-Retriever framework.

The details of the framework can be found here. Notebook with the code is available here

How to use the model

from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch
import torch.nn.functional as F


MODEL_ID = "prhegde/query-product-relevance-model-ecommerce"
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_ID)

query_text = "sofa with ottoman"
prod_text = "daryl 82 '' wide reversible sofa & chaise with ottoman"

tok_output = tokenizer(query_text, prod_text, padding='max_length',
                      max_length=160, truncation=True, return_tensors='pt',
                      return_attention_mask=True)

input_ids = tok_output.input_ids
attention_masks = tok_output.attention_mask
token_type_ids = tok_output.token_type_ids

output = model(input_ids, attention_mask = attention_masks,
                  token_type_ids = token_type_ids)
probs = F.softmax(output.logits, dim=-1)
score = probs[0][1].item()
print(score)