add model use information

README.md

@@ -8645,4 +8645,96 @@ model-index:
 ---
 # Contextual Document Embeddings (CDE)
 
-Our new model that naturally integrates "context tokens" into the embedding process. More information coming soon!
+Our new model that naturally integrates "context tokens" into the embedding process. 
+
+# Using `cde-small-v1`
+
+Our embedding model needs to be used in *two stages*. The first stage is to gather some dataset information by embedding a subset of the corpus using our "first-stage" model. The second stage is to actually embed queries and documents, conditioning on the corpus information from the first stage. Note that we can do the first stage part offline and only use the second-stage weights at inference time.
+
+
+#### Note on prefixes
+
+*Nota bene*: Like all state-of-the-art embedding models, our model was trained with task-specific prefixes. To do retrieval, you can prepend the following strings to queries & documents:
+
+```python
+query_prefix = "search_query: "
+document_prefix = "search_document: "
+```
+
+## First stage
+
+```python
+minicorpus_size = model.config.transductive_corpus_size
+minicorpus_docs = [ ... ] # Put some strings here that are representative of your corpus, for example by calling random.sample(corpus, k=minicorpus_size)
+assert len(minicorpus_docs) == minicorpus_size # You must use exactly this many documents in the minicorpus. You can oversample if your corpus is smaller.
+minicorpus_docs = tokenizer(
+    [document_prefix + doc for doc in minicorpus_docs],
+    truncation=True,
+    padding=True,
+    max_length=512,
+    return_tensors="pt"
+)
+import torch
+from tqdm.autonotebook import tqdm
+
+batch_size = 32
+
+dataset_embeddings = []
+for i in tqdm(range(0, len(minicorpus_docs["input_ids"]), batch_size)):
+    minicorpus_docs_batch = {k: v[i:i+batch_size] for k,v in minicorpus_docs.items()}
+    with torch.no_grad():
+        dataset_embeddings.append(
+            model.first_stage_model(**minicorpus_docs_batch)
+        )
+
+dataset_embeddings = torch.cat(dataset_embeddings)
+
+## Running the second stage
+
+Now that we have obtained "dataset embeddings" we can embed documents and queries like normal. Remember to use the document prefix for documents:
+```python
+docs = tokenizer(
+    [document_prefix + doc for doc in docs],
+    truncation=True,
+    padding=True,
+    max_length=512,
+    return_tensors="pt"
+).to(device)
+
+with torch.no_grad():
+  doc_embeddings = model.second_stage_model(
+      input_ids=docs["input_ids"],
+      attention_mask=docs["attention_mask"],
+      dataset_embeddings=dataset_embeddings,
+  )
+doc_embeddings /= doc_embeddings.norm(p=2, dim=1, keepdim=True)
+```
+
+and the query prefix for queries:
+```python
+queries = queries.select(range(16))["text"]
+queries = tokenizer(
+    [query_prefix + query for query in queries],
+    truncation=True,
+    padding=True,
+    max_length=512,
+    return_tensors="pt"
+).to(device)
+
+with torch.no_grad():
+  query_embeddings = model.second_stage_model(
+      input_ids=queries["input_ids"],
+      attention_mask=queries["attention_mask"],
+      dataset_embeddings=dataset_embeddings,
+  )
+query_embeddings /= query_embeddings.norm(p=2, dim=1, keepdim=True)
+```
+
+these embeddings can be compared using dot product, since they're normalized.
+
+### What if I don't know what my corpus will be ahead of time?
+
+### Colab demo
+
+We've set up a short demo in a Colab notebook showing how you might use our model:
+[Try our model in Colab:](https://colab.research.google.com/drive/1r8xwbp7_ySL9lP-ve4XMJAHjidB9UkbL?usp=sharing)