readme: enrich How-to-run with the original LFM2-ColBERT-350M prose + step-comments

README.md

@@ -90,18 +90,22 @@ We recommend LFM2.5-Embedding-350M and LFM2.5-ColBERT-350M for short-context ret
 
 ## 🏃 How to run
 
-Install PyLate:
+First, install the PyLate and transformers libraries:
 
 ```bash
 pip install -U pylate
 ```
 
+> **Note on `trust_remote_code`.** Loading requires `trust_remote_code=True` so the repo's `modeling_lfm2_bidirectional.py` can replace the causal attention mask and short-conv padding with their non-causal equivalents. Without it the model will silently use causal attention and produce poor retrieval embeddings.
+
 ### Retrieval
 
 Use this model with PyLate to index and retrieve documents. The index uses [FastPLAID](https://github.com/lightonai/fast-plaid) for efficient similarity search.
 
 #### Indexing documents
 
+Load LFM2.5-ColBERT-350M and initialize the PLAID index, then encode and index your documents:
+
 ```python
 from pylate import indexes, models, retrieve
 
@@ -116,7 +120,7 @@ model.tokenizer.pad_token = model.tokenizer.eos_token
 index = indexes.PLAID(
     index_folder="pylate-index",
     index_name="index",
-    override=True,  # Overwrites any existing index with the same name
+    override=True,  # This overwrites the existing index if any
 )
 
 # Step 3: Encode the documents
@@ -126,68 +130,86 @@ documents = ["document 1 text", "document 2 text", "document 3 text"]
 documents_embeddings = model.encode(
     documents,
     batch_size=32,
-    is_query=False,
+    is_query=False,  # Ensure that it is set to False to indicate that these are documents, not queries
     show_progress_bar=True,
 )
 
-# Step 4: Add document embeddings to the index
+# Step 4: Add document embeddings to the index by providing embeddings and corresponding ids
 index.add_documents(
     documents_ids=documents_ids,
     documents_embeddings=documents_embeddings,
 )
 ```
 
-You don't have to rebuild the index every time. Reload an existing index by instantiating it with the same folder/name and `override=False`:
+Note that you do not have to recreate the index and encode the documents every time. Once you have created an index and added the documents, you can re-use the index later by loading it:
 
 ```python
+# To load an index, simply instantiate it with the correct folder/name and without overriding it
 index = indexes.PLAID(
     index_folder="pylate-index",
     index_name="index",
 )
 ```
 
-#### Retrieving top-k documents
+#### Retrieving top-k documents for queries
+
+Once the documents are indexed, you can retrieve the top-k most relevant documents for a given set of queries. To do so, initialize the ColBERT retriever with the index you want to search in, encode the queries, and then retrieve the top-k documents to get the top matches ids and relevance scores:
 
 ```python
+# Step 1: Initialize the ColBERT retriever
 retriever = retrieve.ColBERT(index=index)
 
+# Step 2: Encode the queries
 queries_embeddings = model.encode(
     ["query for document 3", "query for document 1"],
     batch_size=32,
-    is_query=True,
+    is_query=True,  # Ensure that it is set to True to indicate that these are queries
     show_progress_bar=True,
 )
 
+# Step 3: Retrieve top-k documents
 scores = retriever.retrieve(
     queries_embeddings=queries_embeddings,
-    k=10,
+    k=10,  # Retrieve the top 10 matches for each query
 )
 ```
 
 ### Reranking
 
-For reranking on top of a first-stage pipeline, no index is needed:
+If you only want to use LFM2.5-ColBERT-350M to perform reranking on top of your first-stage retrieval pipeline without building an index, you can simply use the `rank` function and pass the queries and documents to rerank:
 
 ```python
 from pylate import rank, models
 
-model = models.ColBERT(
-    model_name_or_path="LiquidAI/LFM2.5-ColBERT-350M",
-    trust_remote_code=True,
-)
+queries = [
+    "query A",
+    "query B",
+]
 
-queries = ["query A", "query B"]
 documents = [
     ["document A", "document B"],
     ["document 1", "document C", "document B"],
 ]
+
 documents_ids = [
     [1, 2],
     [1, 3, 2],
 ]
 
-queries_embeddings = model.encode(queries, is_query=True)
-documents_embeddings = model.encode(documents, is_query=False)
+model = models.ColBERT(
+    model_name_or_path="LiquidAI/LFM2.5-ColBERT-350M",
+    trust_remote_code=True,
+)
+
+queries_embeddings = model.encode(
+    queries,
+    is_query=True,
+)
+
+documents_embeddings = model.encode(
+    documents,
+    is_query=False,
+)
 
 reranked_documents = rank.rerank(
     documents_ids=documents_ids,
@@ -196,8 +218,6 @@ reranked_documents = rank.rerank(
 )
 ```
 
-> **Note on `trust_remote_code`.** Loading requires `trust_remote_code=True` so the repo's `modeling_lfm2_bidirectional.py` can replace the causal attention mask and short-conv padding with their non-causal equivalents. Without it the model will silently use causal attention and produce poor retrieval embeddings.
-
 ## 📈 Performance
 
 Evaluated on two open multilingual retrieval benchmarks, with queries and documents in the same language. Bold = best per column.