<jupyter_start><jupyter_text>Using Our TAS-B Bert_Dot (or BERT Dense Retrieval) CheckpointWe provide a fully retrieval trained (with topic aware and balanced margin sampling: TAS-B) DistilBert-based instance on the HuggingFace model hub here: https://huggingface.co/sebastian-hofstaetter/distilbert-dot-tas_b-b256-msmarcoThis instance can be used to **re-rank a candidate set** or **directly for a vector index based dense retrieval**. The architecure is a 6-layer DistilBERT, without architecture additions or modifications (we only change the weights during training) - to receive a query/passage representation we pool the CLS vector. If you want to know more about our efficient batch composition procedure and dual supervision for dense retrieval training, check out our paper: https://arxiv.org/abs/2104.06967 🎉This notebook gives you a minimal usage example of downloading our Bert_Dot checkpoint to encode passages and queries to create a dot-product based score of their relevance. ---Let's get started by installing the awesome *transformers* library from HuggingFace:<jupyter_code>pip install transformers<jupyter_output><empty_output><jupyter_text>The next step is to download our checkpoint and initialize the tokenizer and models:<jupyter_code>from transformers import AutoTokenizer, AutoModel

# you can switch the model to the original "distilbert-base-uncased" to see that the usage example then breaks and the score ordering is reversed :O
#pre_trained_model_name = "distilbert-base-uncased"
pre_trained_model_name = "sebastian-hofstaetter/distilbert-dot-tas_b-b256-msmarco"

tokenizer = AutoTokenizer.from_pretrained(pre_trained_model_name) 
bert_model = AutoModel.from_pretrained(pre_trained_model_name)<jupyter_output>Downloading: 100%|██████████| 504/504 [00:00<00:00, 128kB/s]
Downloading: 100%|██████████| 262k/262k [00:02<00:00, 91.1kB/s]
Downloading: 100%|██████████| 112/112 [00:00<00:00, 18.7kB/s]
Downloading: 100%|██████████| 306/306 [00:00<00:00, 61.9kB/s]
Downloading: 100%|██████████| 265M/265M [13:03<00:00, 339kB/s]<jupyter_text>Now we are ready to use the model to encode two sample passages and a query:<jupyter_code># our relevant example
passage1_input = tokenizer("We are very happy to show you the 🤗 Transformers library for pre-trained language models. We are helping the community work together towards the goal of advancing NLP 🔥.",return_tensors="pt")
# a non-relevant example
passage2_input = tokenizer("Hmm I don't like this new movie about transformers that i got from my local library. Those transformers are robots?",return_tensors="pt")
# the user query -> which should give us a better score for the first passage
query_input = tokenizer("what is the transformers library",return_tensors="pt")

print("Passage 1 Tokenized:",tokenizer.convert_ids_to_tokens(passage1_input["input_ids"][0]))
print("Passage 2 Tokenized:",tokenizer.convert_ids_to_tokens(passage2_input["input_ids"][0]))
print("Query Tokenized:",tokenizer.convert_ids_to_tokens(query_input["input_ids"][0]))

# note how we call the bert model independently between passages and query :)
# [0][:,0,:] pools (or selects) the CLS vector from the full output
passage1_encoded = bert_model(**passage1_input)[0][:,0,:].squeeze(0)
passage2_encoded = bert_model(**passage2_input)[0][:,0,:].squeeze(0)
query_encoded    = bert_model(**query_input)[0][:,0,:].squeeze(0)

print("---")
print("Passage Encoded Shape:",passage1_encoded.shape)
print("Query Encoded Shape:",query_encoded.shape)<jupyter_output>Passage 1 Tokenized: ['[CLS]', 'we', 'are', 'very', 'happy', 'to', 'show', 'you', 'the', '[UNK]', 'transformers', 'library', 'for', 'pre', '-', 'trained', 'language', 'models', '.', 'we', 'are', 'helping', 'the', 'community', 'work', 'together', 'towards', 'the', 'goal', 'of', 'advancing', 'nl', '##p', '[UNK]', '.', '[SEP]']
Passage 2 Tokenized: ['[CLS]', 'hmm', 'i', 'don', "'", 't', 'like', 'this', 'new', 'movie', 'about', 'transformers', 'that', 'i', 'got', 'from', 'my', 'local', 'library', '.', 'those', 'transformers', 'are', 'robots', '?', '[SEP]']
Query Tokenized: ['[CLS]', 'what', 'is', 'the', 'transformers', 'library', '[SEP]']
---
Passage Encoded Shape: torch.Size([768])
Query Encoded Shape: torch.Size([768])<jupyter_text>Now that we have our encoded vectors, we can generate the score with a simple dot product! 

(This can be offloaded to a vector indexing library like Faiss)<jupyter_code>score_for_p1 = query_encoded.dot(passage1_encoded)
print("Score passage 1 <-> query: ",float(score_for_p1))

score_for_p2 = query_encoded.dot(passage2_encoded)
print("Score passage 2 <-> query: ",float(score_for_p2))<jupyter_output>Score passage 1 <-> query:  103.33966064453125
Score passage 2 <-> query:  97.69734954833984
{'fraction_non_alphanumeric': 0.14638261970048513, 'fraction_numerical': 0.040497785277367644, 'mean_word_length': 6.761047463175123, 'pattern_counts': {'":': 0, '<': 23, '<?xml version=': 0, '>': 19, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 7, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '370216', 'license': 'Apache-2.0', 'n_tokens_mistral': 1728, 'n_tokens_neox': 1536, 'n_words': 477, 'path': 'minimal_bert_dot_usage_example.ipynb', 'repo_name': 'sebastian-hofstaetter/tas-balanced-dense-retrieval', 'size': '4741'}