P-MADRAL is a BERT-sized multi-aspects dense retriever initialized from
BERT public checkpoint,
further pre-trained on e-commerce product data, and fine-tuned on product search retrieval task on the
Amazon ESCI dataset.
It uses a symmetric encoder architecture, with a single shared encoder for both queries and products.
The similarity function is
dot product.
P-MADRAL 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 Amazon ESCI as an example.
6queries = [
7 "iphone 11 pro max case"
8]
9products = [
10 "OtterBox Symmetry Series Case for iPhone 11 Pro Max - Black [...]",
11 "Camera Lens Protector for iPhone 11 Pro/Pro Max, Tempered Glass 9H Hardness Anti-Scratch Camera Screen Protective [...]"
12]
13
14# Load the tokenizer and model.
15tokenizer = transformers.AutoTokenizer.from_pretrained("J-MADRAL/P-MADRAL")
16model = modeling.BiEncoderModel.from_pretrained("J-MADRAL/P-MADRAL")
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")
25p_input = tokenizer(products,
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
34p_emb = model(**p_input).pooled_output
35
36# Compute similarity scores, using dot product similarity.
37scores = torch.matmul(q_emb, p_emb.transpose(0, 1))
