J-MADRAL/R-BiBERT

R-BiBERT is a BERT-sized dense retriever initialized from BERT public checkpoint, further pre-trained on e-commerce review data, and fine-tuned on review search retrieval task on the Search ESCI dataset. It uses a symmetric encoder architecture, with a single shared encoder for both queries and products. The similarity function is dot product.

Paper and Repository

R-BiBERT has been described in the Multi-Aspect Joint Retrieval for E-Commerce: Bridging Product Catalogs and Customer Reviews paper. The associated GitHub repository is available at https://anonymous.4open.science/r/J-MADRAL-C4CC.

Usage (HuggingFace Transformers)

Using the model directly in HuggingFace transformers requires additional code available in the repository.

import modeling
import torch
import transformers

# We use a training query from Search ESCI as an example.
queries = [
    "cotton summer dress care instructions"
]
reviews = [
    "Cute, cool and comfy summer dress [...] Hand wash and line drys easily, material is crinkly so no ironing needed. [...]",
    "Excellent machine JET J-2530 15-Inch 3/4-Horsepower Bench Drill Press. Two common Amazon reviewer complaints about higher-end drill presses [...]"
]

# Load the tokenizer and model.
tokenizer = transformers.AutoTokenizer.from_pretrained("J-MADRAL/R-BiBERT")
model = modeling.BiEncoderModel.from_pretrained("J-MADRAL/R-BiBERT")

# Tokenize the input data.
q_input = tokenizer(queries,
                    add_special_tokens=True,
                    truncation=True,
                    padding=True,
                    max_length=128,
                    return_tensors="pt")
r_input = tokenizer(reviews,
                    add_special_tokens=True,
                    truncation=True,
                    padding=True,
                    max_length=128,
                    return_tensors="pt")

# Compute embeddings: take the "pooled_output".
q_emb = model(**q_input).pooled_output
r_emb = model(**r_input).pooled_output

# Compute similarity scores, using dot product similarity.
scores = torch.matmul(q_emb, r_emb.transpose(0, 1))

Training Hyperparameters

Training Stage	Num. Epochs	Learning Rate	AP Scaling Factor	Max Num Tokens	Batch Size	Num Negatives
Pre-training	20	1e-4	0.10	128	64	---
Fine-tuning	20	5e-6	0.05	128	64	7

The data used for fine-tuning is available at https://huggingface.co/datasets/J-MADRAL/TrainingData.

Evaluation Results

Search ESCI

Model	R@100	R@500	MRR	nDCG@10	nDCG@50
BM25	0.5875	0.7288	0.2539	0.2766	0.3101
DRAGON	0.5451	0.6751	0.2347	0.2567	0.2873
R-BiBERT	0.5972	0.7350	0.2602	0.2855	0.3176
R-MADRAL	0.6405	0.7626	0.2944	0.3215	0.3541
J-BiBERT	0.6300	0.7593	0.2879	0.3140	0.3474
J-MADRAL	0.6488	0.7729	0.3007	0.3281	0.3611