reorg readme

README.md

@@ -8662,6 +8662,111 @@ Our new model that naturally integrates "context tokens" into the embedding proc
 
 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.
 
+</details>
+
+## With Transformers
+
+<details>
+<summary>Click to learn how to use cde-small-v1 with Transformers</summary>
+
+### Loading the model
+
+Our model can be loaded using `transformers` out-of-the-box with "trust remote code" enabled. We use the default BERT uncased tokenizer:
+```python
+import transformers
+
+model = transformers.AutoModel.from_pretrained("jxm/cde-small-v1", trust_remote_code=True)
+tokenizer = transformers.AutoTokenizer.from_pretrained("bert-base-uncased")
+```
+
+#### 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"
+).to(model.device)
+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(model.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(model.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.
+
+</details>
+
+### What if I don't know what my corpus will be ahead of time?
+
+If you can't obtain corpus information ahead of time, you still have to pass *something* as the dataset embeddings; our model will work fine in this case, but not quite as well; without corpus information, our model performance drops from 65.0 to 63.8 on MTEB. We provide [some random strings](https://huggingface.co/jxm/cde-small-v1/resolve/main/random_strings.txt) that worked well for us that can be used as a substitute for corpus sampling.
+
+
 ## With Sentence Transformers
 
 <details open="">
@@ -8832,110 +8937,6 @@ Top Document: Foster's Home for Imaginary Friends McCracken conceived the series
 
 </details>
 
-</details>
-
-## With Transformers
-
-<details>
-<summary>Click to learn how to use cde-small-v1 with Transformers</summary>
-
-### Loading the model
-
-Our model can be loaded using `transformers` out-of-the-box with "trust remote code" enabled. We use the default BERT uncased tokenizer:
-```python
-import transformers
-
-model = transformers.AutoModel.from_pretrained("jxm/cde-small-v1", trust_remote_code=True)
-tokenizer = transformers.AutoTokenizer.from_pretrained("bert-base-uncased")
-```
-
-#### 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"
-).to(model.device)
-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(model.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(model.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.
-
-</details>
-
-### What if I don't know what my corpus will be ahead of time?
-
-If you can't obtain corpus information ahead of time, you still have to pass *something* as the dataset embeddings; our model will work fine in this case, but not quite as well; without corpus information, our model performance drops from 65.0 to 63.8 on MTEB. We provide [some random strings](https://huggingface.co/jxm/cde-small-v1/resolve/main/random_strings.txt) that worked well for us that can be used as a substitute for corpus sampling.
-
 ### Colab demo
 
 We've set up a short demo in a Colab notebook showing how you might use our model: