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.
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.
Using the model directly in HuggingFace transformers requires additional code available in the
repository.
1import modeling
2import torch
3import transformers
4
5# We use a training query from Search ESCI as an example.
6queries = [
7 "cotton summer dress care instructions"
8]
9reviews = [
10 "Cute, cool and comfy summer dress [...] Hand wash and line drys easily, material is crinkly so no ironing needed. [...]",
11 "Excellent machine JET J-2530 15-Inch 3/4-Horsepower Bench Drill Press. Two common Amazon reviewer complaints about higher-end drill presses [...]"
12]
13
14# Load the tokenizer and model.
15tokenizer = transformers.AutoTokenizer.from_pretrained("J-MADRAL/R-BiBERT")
16model = modeling.BiEncoderModel.from_pretrained("J-MADRAL/R-BiBERT")
17
18# Tokenize the input data.
19q_input = tokenizer(queries,
20 add_special_tokens=True,
21 truncation=True,
22 padding=True,
23 max_length=128,
24 return_tensors="pt")
25r_input = tokenizer(reviews,
26 add_special_tokens=True,
27 truncation=True,
28 padding=True,
29 max_length=128,
30 return_tensors="pt")
31
32# Compute embeddings: take the "pooled_output".
33q_emb = model(**q_input).pooled_output
34r_emb = model(**r_input).pooled_output
35
36# Compute similarity scores, using dot product similarity.
37scores = torch.matmul(q_emb, r_emb.transpose(0, 1))
