save

lilac/data/cluster_titling.py

@@ -14,6 +14,7 @@
 from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_random_exponential
 
 from ..batch_utils import group_by_sorted_key_iter
+from ..env import env
 from ..schema import (
   Item,
 )
@@ -24,7 +25,6 @@
 )
 from ..tasks import TaskInfo
 from ..utils import chunks, log
-from ..env import env
 
 _TOP_K_CENTRAL_DOCS = 7
 _TOP_K_CENTRAL_TITLES = 20

lilac/embeddings/bge.py

@@ -2,6 +2,7 @@
 import gc
 from typing import TYPE_CHECKING, ClassVar, Optional
 
+import numpy as np
 from typing_extensions import override
 
 from ..utils import log
@@ -67,7 +68,7 @@ def setup(self) -> None:
   def compute(self, docs: list[str]) -> list[Optional[Item]]:
     """Call the embedding function."""
 
-    def _encode(doc: list[str]):
+    def _encode(doc: list[str]) -> list[np.ndarray]:
       # Extract the dense vectors from the model.
       return self._model.encode(doc)['dense_vecs']
 

lilac/embeddings/nomic_embed.py

@@ -2,6 +2,7 @@
 import gc
 from typing import TYPE_CHECKING, ClassVar, Optional
 
+import numpy as np
 from typing_extensions import override
 
 if TYPE_CHECKING:
@@ -65,7 +66,7 @@ def compute(self, docs: list[str]) -> list[Optional[Item]]:
         'Please install it with `pip install "sentence_transformers".'
       )
 
-    def _encode(doc: list[str]):
+    def _encode(doc: list[str]) -> list[np.ndarray]:
       embeddings = self._model.encode(doc, convert_to_tensor=True)
       # Extract the dense vectors from the model.
       embeddings = F.layer_norm(embeddings, normalized_shape=(embeddings.shape[1],))

mypy.ini

@@ -159,3 +159,8 @@ follow_imports = skip
 [mypy-hdbscan.*]
 ignore_missing_imports = True
 follow_imports = skip
+
+
+[mypy-FlagEmbedding.*]
+ignore_missing_imports = True
+follow_imports = skip