colbertv2 docs

docs/api/retrieval_model_clients/ColBERTv2.md

@@ -49,3 +49,81 @@ retrieval_response = colbertv2_wiki17_abstracts('When was the first FIFA World C
 for result in retrieval_response:
  print("Text:", result['text'], "\n")
 ```
+
+# dspy.ColBERTv2RetrieverLocal
+
+This is taken from the official documentation of [Colbertv2](https://github.com/stanford-futuredata/ColBERT/tree/main) following the [paper](https://arxiv.org/abs/2112.01488).
+
+You can install Colbertv2 by the following instructions from [here](https://github.com/stanford-futuredata/ColBERT?tab=readme-ov-file#installation)
+
+### Constructor
+The constructor initializes the ColBERTv2 as a local retriever object. You can initialize a server instance from your ColBERTv2 local instance using the code snippet from [here](https://github.com/stanford-futuredata/ColBERT/blob/main/server.py)
+
+```python
+class ColBERTv2RetrieverLocal:
+ def __init__(
+ self,
+ passages:List[str],
+ colbert_config=None,
+ load_only:bool=False):
+```
+
+**Parameters**
+- `passages` (_List[str]_): List of passages to be indexed
+- `colbert_config` (_ColBERTConfig_, _Optional_): colbert config for building and searching. Defaults to None.
+- `load_only` (_Boolean_): whether to load the index or build and then load. Defaults to False.
+
+The `colbert_config` object is required for ColBERTv2, and it can be imported from `from colbert.infra.config import ColBERTConfig`. You can find the descriptions of config attributes from [here](https://github.com/stanford-futuredata/ColBERT/blob/main/colbert/infra/config/settings.py)
+
+### Methods
+
+#### `forward(self, query:str, k:int, **kwargs) -> Union[list[str], list[dotdict]]`
+
+It retrieves relevant passages from the index based on the query. If you already have a local index, then you can pass the `load_only` flag as `True` and change the `index` attribute of ColBERTConfig to the local path. Also, make sure to change the `checkpoint` attribute of ColBERTConfig to the embedding model that you used to build the index.
+
+**Parameters:**
+- `query` (_str_): Query string used for retrieval.
+- `k` (_int_, _optional_): Number of passages to retrieve. Defaults to 7
+
+It returns a `Prediction` object for each query
+
+```python
+Prediction(
+ pid=[33, 6, 47, 74, 48],
+ passages=['No pain, no gain.', 'The best things in life are free.', 'Out of sight, out of mind.', 'To be or not to be, that is the question.', 'Patience is a virtue.']
+)
+```
+# dspy.ColBERTv2RerankerLocal
+
+You can also use ColBERTv2 as a reranker in DSPy.
+
+### Constructor
+
+```python
+class ColBERTv2RerankerLocal:
+
+ def __init__(
+ self,
+ colbert_config=None,
+ checkpoint:str='bert-base-uncased'):
+```
+
+**Parameters**
+- `colbert_config` (_ColBERTConfig_, _Optional_): colbert config for building and searching. Defaults to None.
+- `checkpoint` (_str_): Embedding model for embeddings the documents and query
+
+### Methods
+#### `forward(self,query:str,passages:List[str])`
+
+Based on a query and list of passages, it reranks the passages and returns the scores along with the passages ordered in descending order based on the similarity scores.
+
+**Parameters:**
+- `query` (_str_): Query string used for reranking.
+- `passages` (_List[str]_): List of passages to be reranked
+
+It returns the similarity scores array and you can link it to the passages by
+
+```python
+for idx in np.argsort(scores_arr)[::-1]:
+ print(f"Passage = {passages[idx]} --> Score = {scores_arr[idx]}")
+```

dsp/modules/colbertv2.py

@@ -164,7 +164,7 @@ def __init__(self,colbert_config=None,checkpoint:str='bert-base-uncased'):
  checkpoint_name (str, optional): checkpoint for embeddings. Defaults to 'bert-base-uncased'.
  """
  self.colbert_config = colbert_config
- self.checkpoint_name = checkpoint
+ self.checkpoint = checkpoint
  self.colbert_config.checkpoint = checkpoint
 
  def __call__(self, *args: Any, **kwargs: Any) -> Any:
@@ -184,7 +184,7 @@ def forward(self,query:str,passages:List[str]=[]):
  query_ids,query_masks = query_tokenizer.tensorize([query])
  doc_ids,doc_masks = doc_tokenizer.tensorize(passages)
 
- col = ColBERT(self.checkpoint_name,self.colbert_config) 
+ col = ColBERT(self.checkpoint,self.colbert_config) 
  Q = col.query(query_ids,query_masks)
  DOC_IDS,DOC_MASKS = col.doc(doc_ids,doc_masks,keep_dims='return_mask')
  Q_duplicated = Q.repeat_interleave(len(passages), dim=0).contiguous()