Add mypy for type checking

Author: cbornet

pyproject.toml

@@ -47,13 +47,26 @@ log_cli_level = "INFO"
 log_cli_format = "%(asctime)s [%(levelname)8s] %(message)s (%(filename)s:%(lineno)s)"
 log_cli_date_format = "%Y-%m-%d %H:%M:%S"
 
+[tool.mypy]
+disallow_any_generics = true
+disallow_incomplete_defs = true
+disallow_untyped_calls = true
+disallow_untyped_decorators = true
+disallow_untyped_defs = true
+follow_imports = "normal"
+ignore_missing_imports = true
+no_implicit_reexport = true
+show_error_codes = true
+show_error_context = true
+strict_equality = true
+strict_optional = true
+warn_redundant_casts = true
+warn_return_any = true
+warn_unused_ignores = true
 
 [build-system]
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"
 
-
-
-
 [tool.poetry.group.dev.dependencies]
 yamllint = "^1.34.0"

ragstack/colbert/__init__.py

@@ -5,15 +5,15 @@
 from .vector_store import ColBERTVectorStore
 from .constant import DEFAULT_COLBERT_MODEL, DEFAULT_COLBERT_DIM
 
-__all__ = (
- ColbertTokenEmbeddings,
- CassandraColBERTVectorStore,
- ColbertCassandraRetriever,
- max_similarity_torch,
- PerTokenEmbeddings,
- PassageEmbeddings,
- TokenEmbeddings,
- ColBERTVectorStore,
- DEFAULT_COLBERT_MODEL,
- DEFAULT_COLBERT_DIM,
-)
+__all__ = [
+ "ColbertTokenEmbeddings",
+ "CassandraColBERTVectorStore",
+ "ColbertCassandraRetriever",
+ "max_similarity_torch",
+ "PerTokenEmbeddings",
+ "PassageEmbeddings",
+ "TokenEmbeddings",
+ "ColBERTVectorStore",
+ "DEFAULT_COLBERT_MODEL",
+ "DEFAULT_COLBERT_DIM",
+]

ragstack/colbert/cassandra_retriever.py

@@ -1,10 +1,12 @@
-from typing import List
+from typing import List, Set, Tuple, Any, Dict
+
+from cassandra.cluster import ResponseFuture
 
 from .colbert_embedding import ColbertTokenEmbeddings
 
 from .cassandra_store import CassandraColBERTVectorStore
 import logging
-from torch import tensor
+from torch import tensor, Tensor
 import torch
 import math
 
@@ -43,7 +45,9 @@ def max_similarity_numpy_based(query_vector, embedding_list):
 
 # this torch based max similary has the best performance.
 # it is at least 20 times faster than dot product operator and numpy based implementation CuDA and CPU
-def max_similarity_torch(query_vector, embedding_list, is_cuda: bool = False):
+def max_similarity_torch(
+ query_vector: Tensor, embedding_list: List[Tensor], is_cuda: bool = False
+) -> Tensor:
  """
  Calculate the maximum similarity (dot product) between a query vector and a list of embedding vectors,
  optimized for performance using PyTorch for GPU acceleration.
@@ -59,12 +63,12 @@ def max_similarity_torch(query_vector, embedding_list, is_cuda: bool = False):
  # stacks the list of embedding tensors into a single tensor
  if is_cuda:
  query_vector = query_vector.to("cuda")
- embedding_list = torch.stack(embedding_list).to("cuda")
+ _embedding_list = torch.stack(embedding_list).to("cuda")
  else:
- embedding_list = torch.stack(embedding_list)
+ _embedding_list = torch.stack(embedding_list)
 
  # Calculate the dot products in a vectorized manner on the GPU
- sims = torch.matmul(embedding_list, query_vector)
+ sims = torch.matmul(_embedding_list, query_vector)
 
  # Find the maximum similarity (dot product) value
  max_sim = torch.max(sims)
@@ -90,11 +94,11 @@ def __init__(
  self.colbert_embeddings = colbert_embeddings
  self.is_cuda = torch.cuda.is_available()
 
- def close(self):
+ def close(self) -> None:
  pass
 
  def retrieve(
- self, query: str, k: int = 10, query_maxlen: int = 64, **kwargs
+ self, query: str, k: int = 10, query_maxlen: int = 64, **kwargs: Any
  ) -> List[Document]:
  #
  # if the query has fewer than a predefined number of tokens Nq,
@@ -109,7 +113,7 @@ def retrieve(
  logging.debug(f"query length {len(query)} embeddings top_k: {top_k}")
 
  # find the most relevant documents
- docparts = set()
+ docparts: Set[Tuple[Any, Any]] = set()
  doc_futures = []
  for qv in query_encodings:
  # per token based retrieval
@@ -146,17 +150,17 @@ def retrieve(
  docs_by_score = sorted(scores, key=scores.get, reverse=True)[:k]
 
  # query the doc body
- doc_futures = {}
+ doc_futures2: Dict[Tuple[Any, Any], ResponseFuture] = {}
  for title, part in docs_by_score:
  future = self.vector_store.session.execute_async(
  self.vector_store.query_part_by_pk_stmt, [title, part]
  )
- doc_futures[(title, part)] = future
+ doc_futures2[(title, part)] = future
 
  answers: List[Document] = []
  rank = 1
  for title, part in docs_by_score:
- rs = doc_futures[(title, part)].result()
+ rs = doc_futures2[(title, part)].result()
  score = scores[(title, part)]
  answers.append(
  Document(title=title, score=score.item(), rank=rank, body=rs.one().body)

ragstack/colbert/cassandra_store.py

@@ -65,7 +65,7 @@ def __init__(self, session: Session, keyspace: str, table_name: str):
  """
  )
 
- def __create_tables(self):
+ def __create_tables(self) -> None:
  self.session.execute(
  f"""
  CREATE TABLE IF NOT EXISTS {self.full_table_name} (
@@ -112,10 +112,10 @@ def put_document(
  ) -> None:
  return self.insert_colbert_embeddings_chunks(embeddings, delete_existed_passage)
 
- def delete_documents(self, titles: List[str]):
+ def delete_documents(self, titles: List[str]) -> None:
  execute_concurrent_with_args(
  self.session, self.delete_part_by_title_stmt, [(t,) for t in titles]
  )
 
- def close(self):
+ def close(self) -> None:
  pass

ragstack/colbert/colbert_embedding.py

@@ -106,10 +106,6 @@ def __init__(
  query_maxlen=query_maxlen,
  gpus=total_visible_gpus,
  )
- self.__doc_maxlen = doc_maxlen
- self.__nbits = nbits
- self.__kmeans_niters = kmeans_niters
- self.__nranks = nranks
  logging.info("creating checkpoint")
  self.checkpoint = Checkpoint(
  self.colbert_config.checkpoint, colbert_config=self.colbert_config
@@ -148,7 +144,7 @@ def encode_queries(
  # the length does not grow or shrink despite the number of tokens in the query
  # we continue to use the same term to align with ColBERT documentation/library
  query_maxlen: int = -1,
- ):
+ ) -> Tensor:
  queries = query if isinstance(query, list) else [query]
  bsize = 128 if len(queries) > 128 else None
 
@@ -179,7 +175,7 @@ def encode_query(
  query: str,
  full_length_search: bool = False,
  query_maxlen: int = 32,
- ):
+ ) -> Tensor:
  queries = self.encode_queries(
  query, full_length_search, query_maxlen=query_maxlen
  )

ragstack/colbert/langchain/__init__.py

@@ -1,3 +1,3 @@
 from .retriever import ColBERTVectorStoreLangChainRetriever
 
-__all__ = (ColBERTVectorStoreLangChainRetriever,)
+__all__ = ["ColBERTVectorStoreLangChainRetriever"]

ragstack/colbert/langchain/retriever.py

@@ -21,6 +21,7 @@ class ColBERTVectorStoreLangChainRetriever(BaseRetriever):
  qa = RetrievalQA.from_chain_type(llm=llm, chain_type="stuff", retriever=retriever)
  qa.run("what happened on June 4th?")
  """
+
  retriever: ColBERTVectorStoreRetriever = Field(default=None)
  kwargs: dict = {}
  k: int = 10

ragstack/colbert/token_embedding.py

@@ -2,7 +2,7 @@
 # this is a base class for ColBERT per token based embedding
 
 from abc import ABC, abstractmethod
-from typing import List
+from typing import List, Optional
 from .constant import DEFAULT_COLBERT_DIM, DEFAULT_COLBERT_MODEL
 import uuid
 
@@ -14,7 +14,7 @@ def __init__(
  self,
  id: int,
  part: int,
- parent_id: uuid.UUID = None,
+ parent_id: Optional[uuid.UUID] = None,
  title: str = "",
  ):
  self.id = id
@@ -23,19 +23,19 @@ def __init__(
  self.title = title
  self.part = part
 
- def add_embeddings(self, embeddings: List[float]):
+ def add_embeddings(self, embeddings: List[float]) -> None:
  self.__embeddings = embeddings
 
  def get_embeddings(self) -> List[float]:
  return self.__embeddings
 
- def id(self):
+ def id(self) -> int:
  return self.id
 
- def parent_id(self):
+ def parent_id(self) -> Optional[uuid.UUID]:
  return self.parent_id
 
- def part(self):
+ def part(self) -> int:
  return self.part
 
 
@@ -50,7 +50,7 @@ def __init__(
  text: str,
  title: str = "",
  part: int = 0,
- id: uuid.UUID = None,
+ id: Optional[uuid.UUID] = None,
  model: str = DEFAULT_COLBERT_MODEL,
  dim: int = DEFAULT_COLBERT_DIM,
  ):
@@ -65,31 +65,31 @@ def __init__(
  self.__title = title
  self.__part = part
 
- def model(self):
+ def model(self) -> str:
  return self.__model
 
- def dim(self):
+ def dim(self) -> int:
  return self.__dim
 
- def token_size(self):
+ def token_size(self) -> int:
  return len(self.token_ids)
 
- def title(self):
+ def title(self) -> str:
  return self.__title
 
- def __len__(self):
+ def __len__(self) -> int:
  return len(self.embeddings)
 
- def id(self):
+ def id(self) -> uuid.UUID:
  return self.__id
 
- def part(self):
+ def part(self) -> int:
  return self.__part
 
- def add_token_embeddings(self, token_embeddings: PerTokenEmbeddings):
+ def add_token_embeddings(self, token_embeddings: PerTokenEmbeddings) -> None:
  self.__token_embeddings.append(token_embeddings)
 
- def get_token_embeddings(self, token_id: int) -> PerTokenEmbeddings:
+ def get_token_embeddings(self, token_id: int) -> Optional[PerTokenEmbeddings]:
  for token in self.__token_embeddings:
  if token.token_id == token_id:
  return token
@@ -98,7 +98,7 @@ def get_token_embeddings(self, token_id: int) -> PerTokenEmbeddings:
  def get_all_token_embeddings(self) -> List[PerTokenEmbeddings]:
  return self.__token_embeddings
 
- def get_text(self):
+ def get_text(self) -> str:
  return self.__text
 
 

ragstack/colbert/vector_store.py

@@ -5,24 +5,22 @@
 import dataclasses
 from abc import ABC, abstractmethod
 from numbers import Number
-from typing import List, Optional
+from typing import List, Optional, Any
 
 
 class ColBERTVectorStore(ABC):
  """Interface for a vector store."""
 
  @abstractmethod
- def close(self):
+ def close(self) -> None:
  """Close the store."""
- pass
 
  @abstractmethod
- def put_document(self, document: str, metadata: dict):
+ def put_document(self, document: str, metadata: dict) -> None:
  """Put a document into the store."""
- pass
 
  @abstractmethod
- def delete_documents(self, titles: List[str]):
+ def delete_documents(self, titles: List[str]) -> None:
  """Delete a document from the store."""
  pass
 
@@ -37,13 +35,11 @@ class Document:
 
 class ColBERTVectorStoreRetriever(ABC):
  @abstractmethod
- def close(self):
+ def close(self) -> None:
  """Close the store."""
- pass
 
  @abstractmethod
  def retrieve(
- self, query: str, k: Optional[int], query_maxlen: Optional[int], **kwargs
+ self, query: str, k: Optional[int], query_maxlen: Optional[int], **kwargs: Any
  ) -> List[Document]:
  """Retrieve documents from the store"""
- pass