Update tests

README.md

@@ -47,7 +47,19 @@ from transformers import AutoModelForMaskedLM, AutoTokenizer
 model = model.Splade(
     model=AutoModelForMaskedLM.from_pretrained("distilbert-base-uncased").to(device),
     tokenizer=AutoTokenizer.from_pretrained("distilbert-base-uncased"),
-    device=device
+    device=device,
+)
+```
+
+```python
+from transformers import AutoModelForMaskedLM, AutoTokenizer
+
+model = model.SparsEmbed(
+    model=AutoModelForMaskedLM.from_pretrained("distilbert-base-uncased").to(device),
+    tokenizer=AutoTokenizer.from_pretrained("distilbert-base-uncased"),
+    embedding_size=64,
+    k_tokens=96,
+    device=device,
 )
 ```
 

sparsembed/losses/flops.py

@@ -11,6 +11,9 @@ class Flops(torch.nn.Module):
     >>> from transformers import AutoModelForMaskedLM, AutoTokenizer
     >>> from sparsembed import model, utils, losses
     >>> from pprint import pprint as print
+    >>> import torch
+
+    >>> _ = torch.manual_seed(42)
 
     >>> device = "mps"
 
@@ -37,7 +40,7 @@ class Flops(torch.nn.Module):
     ...     positive_activations=positive_activations["sparse_activations"],
     ...     negative_activations=negative_activations["sparse_activations"],
     ... )
-    tensor(1838.6599, device='mps:0')
+    tensor(1880.5656, device='mps:0')
 
     References
     ----------

sparsembed/model/sparsembed.py

@@ -17,10 +17,8 @@ class SparsEmbed(Splade):
         HuggingFace Tokenizer.
     model
         HuggingFace AutoModelForMaskedLM.
-    k_query
-        Number of activated terms to retrieve for queries.
-    k_documents
-        Number of activated terms to retrieve for documents.
+    k_tokens
+        Number of activated terms to retrieve.
     embedding_size
         Size of the embeddings in output of SparsEmbed model.
 

sparsembed/train/train_sparsembed.py

@@ -41,6 +41,8 @@ def train_sparsembed(
     >>> from sparsembed import model, utils, train
     >>> import torch
 
+    >>> _ = torch.manual_seed(42)
+
     >>> device = "mps"
 
     >>> model = model.SparsEmbed(
@@ -73,7 +75,9 @@ def train_sparsembed(
     ...         in_batch_negatives=True,
     ...     )
 
-    >>> {'dense': tensor(4.9316, device='mps:0', grad_fn=<MeanBackward0>), 'ranking': tensor(3456.6538, device='mps:0', grad_fn=<MeanBackward0>), 'flops': tensor(796.2637, device='mps:0', grad_fn=<SumBackward1>)}
+    >>> loss
+    {'dense': tensor(19.4581, device='mps:0', grad_fn=<MeanBackward0>), 'sparse': tensor(3475.2351, device='mps:0', grad_fn=<MeanBackward0>), 'flops': tensor(795.3960, device='mps:0', grad_fn=<SumBackward1>)}
+
     """
 
     anchor_activations = model(
@@ -127,6 +131,6 @@ def train_sparsembed(
 
     return {
         "dense": dense_ranking_loss,
-        "ranking": sparse_ranking_loss,
+        "sparse": sparse_ranking_loss,
         "flops": flops_loss,
     }

sparsembed/train/train_splade.py

@@ -38,6 +38,8 @@ def train_splade(
     >>> from sparsembed import model, utils, train
     >>> import torch
 
+    >>> _ = torch.manual_seed(42)
+
     >>> device = "mps"
 
     >>> model = model.Splade(
@@ -71,7 +73,7 @@ def train_splade(
     ...     )
 
     >>> loss
-    {'ranking': tensor(307.2816, device='mps:0', grad_fn=<MeanBackward0>), 'flops': tensor(75.3216, device='mps:0', grad_fn=<SumBackward1>)}
+    {'sparse': tensor(5443.0615, device='mps:0', grad_fn=<MeanBackward0>), 'flops': tensor(1319.7771, device='mps:0', grad_fn=<SumBackward1>)}
 
     """
 
@@ -108,4 +110,4 @@ def train_splade(
     optimizer.step()
     optimizer.zero_grad()
 
-    return {"ranking": ranking_loss, "flops": flops_loss}
+    return {"sparse": ranking_loss, "flops": flops_loss}

sparsembed/utils/evaluate.py

@@ -63,6 +63,9 @@ def evaluate(
 
     >>> from transformers import AutoModelForMaskedLM, AutoTokenizer
     >>> from sparsembed import model, retrieve, utils
+    >>> import torch
+
+    >>> _ = torch.manual_seed(42)
 
     >>> device = "cpu"
 
@@ -87,7 +90,7 @@ def evaluate(
     ...     batch_size=1
     ... )
 
-    >>> utils.evaluate(
+    >>> scores = utils.evaluate(
     ...     retriever=retriever,
     ...     batch_size=1,
     ...     qrels=qrels,
@@ -96,6 +99,9 @@ def evaluate(
     ...     metrics=["map", "ndcg@10", "ndcg@100", "recall@10", "recall@100"]
     ... )
 
+    >>> scores
+    {'map': 0.0016666666666666668, 'ndcg@10': 0.002103099178571525, 'ndcg@100': 0.002103099178571525, 'recall@10': 0.0033333333333333335, 'recall@100': 0.0033333333333333335}
+
     """
     from ranx import Qrels, Run, evaluate