fix typing issues from dataset and embedding models

Author: haok1402

.vscode/settings.json

@@ -9,7 +9,6 @@
   },
   "[python]": {
     "editor.tabSize": 4,
-    "editor.formatOnSave": true,
     "editor.defaultFormatter": "ms-python.black-formatter",
   },
   "[jsonc]": {

source/dataset/msMarco.py

@@ -100,9 +100,12 @@ def docPrefixEmbIter(
         shuffle: bool,
         idxs: List[int],
     ) -> Iterator[Tuple[Tensor, List[str], Tensor]]:
+        """
+        @todo: fix the typing override.
+        """
         embed = embedding()
         idx = 0
-        idxs = deque(sorted(idxs))
+        idxs = deque(sorted(idxs)) # type: ignore
         done = False
         for p in range(4):
             path = Path(DocIterInit.base, f"partition-{p:08d}.parquet")
@@ -111,9 +114,12 @@ def docPrefixEmbIter(
             batches = file.iter_batches(1, columns=["text"])
             for i, part in enumerate(batches):
                 if idx == idxs[0]:
-                    idxs.popleft()
+                    idxs.popleft() # type: ignore
                     txt = part.column("text").to_pylist()
-                    vec, tokens, token_ids = embed.forward_prefix(txt)
+                    """
+                    @todo: add forward_prefix to the Embedding interface.
+                    """
+                    vec, tokens, token_ids = embed.forward_prefix(txt) # type: ignore
                     yield vec, tokens, token_ids.detach().cpu().tolist()
                 idx += 1
                 if len(idxs) == 0:
@@ -396,7 +402,7 @@ def __init__(self) -> None:
         self.base.mkdir(mode=0o770, parents=True, exist_ok=True)
         asyncio.run(self.dispatch())
 
-    async def dispatch(self) -> None:
+    async def dispatch(self):
         # we should have dispatched all tasks at once, but due to progress bar
         # constraints, only one at a time is possible. Otherwise, the progress
         # bar would be globally defined, and may interfere with training logs.

source/dataset/test_msMarco.py

@@ -2,7 +2,7 @@
 from source.embedding.bgeBase import BgeBaseEmbedding
 
 
-def test_didIter_1():
+def test_didIter():
     """
     Test didIter method.
     """
@@ -12,7 +12,7 @@ def test_didIter_1():
     assert all(isinstance(i, int) for i in ids)
 
 
-def test_docIter_1():
+def test_docIter():
     """
     Test docIter method.
     """
@@ -22,7 +22,7 @@ def test_docIter_1():
     assert all(isinstance(d, str) for d in docs)
 
 
-def test_docEmbIter_1():
+def test_docEmbIter():
     """
     Test docEmbIter method.
     """
@@ -31,7 +31,7 @@ def test_docEmbIter_1():
     assert embeddings.shape == (8, BgeBaseEmbedding.size)
 
 
-def test_getDocLen_1():
+def test_getDocLen():
     """
     Test getDocLen method.
     """
@@ -41,119 +41,77 @@ def test_getDocLen_1():
     assert docLen == 8841823
 
 
-def test_qidIter_1():
+def test_qidIter():
     """
     Test qidIter method.
     """
     dataset = MsMarcoDataset()
     qids = next(dataset.qidIter("Train", 8))
     assert isinstance(qids, list) and len(qids) == 8
     assert all(isinstance(q, int) for q in qids)
-
-
-def test_qidIter_2():
-    """
-    Test qidIter method.
-    """
-    dataset = MsMarcoDataset()
     qids = next(dataset.qidIter("Validate", 8))
     assert isinstance(qids, list) and len(qids) == 8
     assert all(isinstance(q, int) for q in qids)
 
 
-def test_qryIter_1():
+def test_qryIter():
     """
     Test qryIter method.
     """
     dataset = MsMarcoDataset()
     qrys = next(dataset.qryIter("Train", 8))
     assert isinstance(qrys, list) and len(qrys) == 8
     assert all(isinstance(q, str) for q in qrys)
-
-
-def test_qryIter_2():
-    """
-    Test qryIter method.
-    """
-    dataset = MsMarcoDataset()
     qrys = next(dataset.qryIter("Validate", 8))
     assert isinstance(qrys, list) and len(qrys) == 8
     assert all(isinstance(q, str) for q in qrys)
 
 
-def test_qryEmbIter_1():
+def test_qryEmbIter():
     """
     Test qryEmbIter method.
     """
     dataset = MsMarcoDataset()
     embeddings = next(dataset.qryEmbIter(BgeBaseEmbedding, "Train", 8, 0, False))
     assert embeddings.shape == (8, BgeBaseEmbedding.size)
-
-
-def test_qryEmbIter_2():
-    """
-    Test qryEmbIter method.
-    """
-    dataset = MsMarcoDataset()
     embeddings = next(dataset.qryEmbIter(BgeBaseEmbedding, "Validate", 8, 0, False))
     assert embeddings.shape == (8, BgeBaseEmbedding.size)
 
 
-def test_getQryLen_1():
+def test_getQryLen():
     """
     Test getQryLen method.
     """
     dataset = MsMarcoDataset()
     qryLen = dataset.getQryLen("Train")
     assert isinstance(qryLen, int)
     assert qryLen == 808731
-
-
-def test_getQryLen_2():
-    """
-    Test getQryLen method.
-    """
-    dataset = MsMarcoDataset()
     qryLen = dataset.getQryLen("Validate")
     assert isinstance(qryLen, int)
     assert qryLen == 101093
 
 
-def test_mixEmbIter_1():
+def test_mixEmbIter():
     """
     Test mixEmbIter method.
     """
     dataset = MsMarcoDataset()
     qry, docs = next(dataset.mixEmbIter(BgeBaseEmbedding, "Train", 32, 8, 0, False))
     assert qry.shape == (8, BgeBaseEmbedding.size)
     assert docs.shape == (8, 32, BgeBaseEmbedding.size)
-
-
-def test_mixEmbIter_2():
-    """
-    Test mixEmbIter method.
-    """
-    dataset = MsMarcoDataset()
     qry, docs = next(dataset.mixEmbIter(BgeBaseEmbedding, "Validate", 32, 8, 0, False))
     assert qry.shape == (8, BgeBaseEmbedding.size)
     assert docs.shape == (8, 32, BgeBaseEmbedding.size)
 
 
-def test_getMixLen_1():
+def test_getMixLen():
     """
     Test getMixLen method.
     """
     dataset = MsMarcoDataset()
     mixLen = dataset.getMixLen("Train")
     assert isinstance(mixLen, int)
     assert mixLen == 808731
-
-
-def test_getMixLen_2():
-    """
-    Test getMixLen method.
-    """
-    dataset = MsMarcoDataset()
     mixLen = dataset.getMixLen("Validate")
     assert isinstance(mixLen, int)
     assert mixLen == 101093

source/embedding/bgeBase.py

@@ -2,7 +2,7 @@
 from torch import Tensor
 import torch.nn as nn
 import torch.nn.functional as F
-from typing import List
+from typing import List, Tuple, Any
 from transformers.models.bert.modeling_bert import BertModel
 from transformers.models.bert.tokenization_bert_fast import BertTokenizerFast
 from transformers.modeling_outputs import BaseModelOutputWithPoolingAndCrossAttentions
@@ -38,7 +38,10 @@ def forward(self, passages: List[str]) -> Tensor:
         return F.normalize(hiddens[:, 0], p=2, dim=1)
 
     @torch.inference_mode()
-    def forward_prefix(self, passages: List[str]) -> Tensor:
+    def forward_prefix(self, passages: List[str]) -> Tuple[Tensor, Any, Any]:
+        """
+        @todo: fix the return type.
+        """
         kwargs = dict(padding=True, truncation=True, return_tensors="pt")
         encoded = self.tokenizer(passages[0], **kwargs)
         input_ids = encoded.input_ids[0]  # Shape: [seq_len]

source/embedding/test_bgeBase.py

@@ -1,7 +1,7 @@
 from source.embedding.bgeBase import BgeBaseEmbedding
 
 
-def test_forward_1():
+def test_forward():
     """
     Test forward method.
     """