Merge pull request #2 from bioinformatics-ua/collection-refactor

Add an AbstractSparseCollection and leave SparseColletion as default imp

converting_from_anserini.py

@@ -13,7 +13,7 @@
 def main(dataset_folder):
 
 
- sparseCSR_collection = SparseCollectionCSR.from_bm25_pyserini_iterator(os.path.join(dataset_folder, "anserini_index"),
+ sparseCSR_collection = SparseCollectionCSR.from_bm25_pyserini_index(os.path.join(dataset_folder, "anserini_index"),
  k1=1.2,
  b=0.75,
  dtype=spare.float32,

spare/collection.py

@@ -1,14 +1,17 @@
 from spare import TYPE
-from spare.utils import get_coo_sparce_GB, get_csr_sparce_GB
+from spare.utils import get_coo_sparce_GB, get_csr_sparce_GB, load_backend, maybe_init
 from tqdm import tqdm
 from spare.metadata import MetaDataDFandDL
 import os
 import jsonpickle
 import shutil
 from spare.weighting_model import WeightingSchemaType, CountingWeightingSchema, BM25WeightingSchema
 
-class SparseCollection:
- ### implements a COO sparse matrix
+
+class AbstractSparseCollection:
+ """
+ A SparseCollection without a concrete implementation
+ """
  def __init__(self, 
  collection_maxsize, 
  text_to_vec=None,
@@ -19,7 +22,6 @@ def __init__(self,
  backend="torch") -> None:
  super().__init__()
 
-
  self.collection_maxsize = collection_maxsize
 
  if isinstance(dtype, int): 
@@ -35,14 +37,10 @@ def __init__(self,
  else:
  self.vec_dim = vec_dim
 
- if backend=="torch":
- from spare.backend_torch import TorchBackend
- self.backend = TorchBackend()
- else:
- RuntimeError("Only torch backend is currently supported")
+ self.backend = load_backend(backend)
 
- self.metadata = metadata()
- self.weighting_schema = weighting_schema()
+ self.metadata = maybe_init(metadata)
+ self.weighting_schema = maybe_init(weighting_schema)
 
  @classmethod
  def from_text_iterator(cls,
@@ -98,18 +96,14 @@ def from_bm25_pyserini_index(cls,
  sparse_collection = cls(index_reader.stats()["documents"], 
  vec_dim=index_reader.stats()["unique_terms"], 
  dtype=dtype, 
- weighting_schema=BM25WeightingSchema,
+ weighting_schema=BM25WeightingSchema(k1=k1, b=b, idf_weighting=None),
  backend=backend, 
  **kwargs)
 
  # probably can skip estimation
  sparse_collection._build_sparse_collection(bm25_pyserini_iterator(index_reader, k1=k1, b=b), 
  max_files_for_estimation=max_files_for_estimation)
-
- sparse_collection.weighting_schema.k1 = k1
- sparse_collection.weighting_schema.b = b
- sparse_collection.weighting_schema.idf_weighting = None
-
+
  return sparse_collection
 
  def transform(self, transform_operator):
@@ -119,10 +113,10 @@ def transform(self, transform_operator):
  raise ValueError("Your current collection weighing schema or metadata is not compatible with the transformation asked.")
 
  def get_sparce_matrix(self):
- return self.backend.create_coo_matrix(*self.sparse_vecs, self.shape)
+ raise NotImplementedError("method get_sparce_matrix was not implemented, if this is an AbstractSparseCollection. Then the behaviour is expected.")
 
  def _correct_transform_operator(self, transform_operator):
- return transform_operator.for_coo()
+ raise NotImplementedError("method _correct_transform_operator was not implemented, if this is an AbstractSparseCollection. Then the behaviour is expected.")
 
  def _build_sparse_collection(self, iterator, max_files_for_estimation):
 
@@ -194,57 +188,17 @@ def _build_sparse_tensor(self, iterator, collection_maxsize, elements_expected,
  return self._slice_sparse_vecs(*sparse_vecs, element_index)
 
  def _slice_sparse_vecs(self, indices, values, element_index):
- # if I slice the indices I will create a non-continguous vector... which is not good
- #indices = self.backend.slice_tensor(indices, (slice(None, None, None), slice(None, element_index, None)))
- #values = self.backend.slice_tensor(values, slice(None, element_index, None))
-
- # for the COO we will pad the documents to +1
- doc_id = self.backend.get_value_from_tensor(indices, (0,element_index-1))+1
- len_pad = values.shape[0]-element_index
- self.backend.assign_data_to_tensor(indices,(0, slice(element_index, None, None)),[doc_id]*len_pad,TYPE.int64)
-
- # update shape to include the padded doc
- self.shape = (self.shape[0]+1, self.shape[1])
- self.metadata.update(doc_id, [], [0]) # pad the metadata
-
- return indices, values
-
+ raise NotImplementedError("method _slice_sparse_vecs was not implemented, if this is an AbstractSparseCollection. Then the behaviour is expected.")
+
  def _init_sparse_vecs(self, elements_expected):
- indices = self.backend.create_zero_tensor((2,elements_expected), TYPE.int64)
- values = self.backend.create_zero_tensor((elements_expected,), self.dtype)
-
- return indices, values
-
- def _sparcify_bow_and_meta_update(self, bow, index_docs):
- py_indices_row = []
- py_indices_col = []
- py_values = []
- for token_index in sorted(bow.keys()):
- py_indices_row.append(index_docs)
- py_indices_col.append(token_index)
- py_values.append(bow[token_index])
-
- self.metadata.update(index_docs, py_indices_col, py_values)
-
- return py_indices_row, py_indices_col, py_values
+ raise NotImplementedError("method _init_sparse_vecs was not implemented, if this is an AbstractSparseCollection. Then the behaviour is expected.")
 
  def _update_sparse_vecs(self, indices, values, bow, element_index, index_docs):
-
- py_indices_row, py_indices_col, py_values = self._sparcify_bow_and_meta_update(bow, index_docs)
-
- self.backend.assign_data_to_tensor(indices, 
- (slice(None, None, None), slice(element_index,element_index+len(py_indices_col), None)),
- [py_indices_row, py_indices_col],
- TYPE.int64)
- self.backend.assign_data_to_tensor(values, 
- slice(element_index,element_index+len(py_indices_col), None),
- py_values,
- self.dtype)
-
- return len(py_indices_col)
+ raise NotImplementedError("method _update_sparse_vecs was not implemented, if this is an AbstractSparseCollection. Then the behaviour is expected.")
+
 
  def get_sparce_matrix_space(self):
- return get_coo_sparce_GB(self.shape, self.density, self.dtype)
+ raise NotImplementedError("method get_sparce_matrix_space was not implemented, if this is an AbstractSparseCollection. Then the behaviour is expected.")
 
  def _get_matrix_estimations(self,
  sampled_bow_list):
@@ -299,8 +253,70 @@ def load_from_file(cls, folder_name, backend="torch"):
  sparse_collection.sparse_vecs = sparse_collection.backend.load_tensors_from_file(os.path.join(folder_name, "tensors.safetensors"))
 
  return sparse_collection
+
+class SparseCollectionCOO(AbstractSparseCollection):
+
+ def get_sparce_matrix(self):
+ return self.backend.create_coo_matrix(*self.sparse_vecs, self.shape)
+
+ def _correct_transform_operator(self, transform_operator):
+ return transform_operator.for_coo()
+
+ def get_sparce_matrix_space(self):
+ return get_coo_sparce_GB(self.shape, self.density, self.dtype)
+
+ def _init_sparse_vecs(self, elements_expected):
+
+ indices = self.backend.create_zero_tensor((2,elements_expected), TYPE.int64)
+ values = self.backend.create_zero_tensor((elements_expected,), self.dtype)
+
+ return indices, values
+
+ def _sparcify_bow_and_meta_update(self, bow, index_docs):
+ py_indices_row = []
+ py_indices_col = []
+ py_values = []
+ for token_index in sorted(bow.keys()):
+ py_indices_row.append(index_docs)
+ py_indices_col.append(token_index)
+ py_values.append(bow[token_index])
+
+ self.metadata.update(index_docs, py_indices_col, py_values)
+
+ return py_indices_row, py_indices_col, py_values
+
+ def _update_sparse_vecs(self, indices, values, bow, element_index, index_docs):
+
+ py_indices_row, py_indices_col, py_values = self._sparcify_bow_and_meta_update(bow, index_docs)
+
+ self.backend.assign_data_to_tensor(indices, 
+ (slice(None, None, None), slice(element_index,element_index+len(py_indices_col), None)),
+ [py_indices_row, py_indices_col],
+ TYPE.int64)
+ self.backend.assign_data_to_tensor(values, 
+ slice(element_index,element_index+len(py_indices_col), None),
+ py_values,
+ self.dtype)
+
+ return len(py_indices_col)
+
+ def _slice_sparse_vecs(self, indices, values, element_index):
+ # if I slice the indices I will create a non-continguous vector... which is not good
+ #indices = self.backend.slice_tensor(indices, (slice(None, None, None), slice(None, element_index, None)))
+ #values = self.backend.slice_tensor(values, slice(None, element_index, None))
 
-class SparseCollectionCSR(SparseCollection):
+ # for the COO we will pad the documents to +1
+ doc_id = self.backend.get_value_from_tensor(indices, (0,element_index-1))+1
+ len_pad = values.shape[0]-element_index
+ self.backend.assign_data_to_tensor(indices,(0, slice(element_index, None, None)),[doc_id]*len_pad,TYPE.int64)
+
+ # update shape to include the padded doc
+ self.shape = (self.shape[0]+1, self.shape[1])
+ self.metadata.update(doc_id, [], [0]) # pad the metadata
+
+ return indices, values
+
+class SparseCollectionCSR(SparseCollectionCOO):
  # sparse collection imlemented with csr
  def __init__(self, 
  *args,
@@ -310,13 +326,13 @@ def __init__(self,
  if isinstance(indices_dtype, int): 
  indices_dtype = TYPE(indices_dtype)
  self.indices_dtype = indices_dtype
-
- def _correct_transform_operator(self, transform_operator):
- return transform_operator.for_csr()
 
  def get_sparce_matrix_space(self):
  return get_csr_sparce_GB(self.shape, self.density, self.dtype, self.indices_dtype)
 
+ def _correct_transform_operator(self, transform_operator):
+ return transform_operator.for_csr()
+
  def get_sparce_matrix(self):
  return self.backend.create_csr_matrix(*self.sparse_vecs, self.shape, self.dtype)
 
@@ -360,4 +376,8 @@ def _slice_sparse_vecs(self, crow_indices, col_indices, values, element_index):
  return crow_indices, col_indices, values
 
  def _get_class_attributes(self):
- return super()._get_class_attributes() | {"indices_dtype":self.indices_dtype.value}
+ return super()._get_class_attributes() | {"indices_dtype":self.indices_dtype.value}
+
+
+# current default implementation
+SparseCollection = SparseCollectionCSR

spare/retriever.py

@@ -49,11 +49,10 @@ def query_transform(query):
  return_scores=return_scores)
 
  # maybe convert index2docID if docID!=index
-
  s_time = time.time()
  converted_indices = []
  for i in range(len(out.ids)):
- q_indices = self.collection.metadata.index2docID[out.ids[i]]#[self.collection.metadata.index2docID[idx] for idx in out.ids[i]]
+ q_indices = self.collection.metadata.index2docID[out.ids[i]]
  converted_indices.append(q_indices)
  print("Time to convert docs ids",time.time()-s_time )
  out.ids = np.array(converted_indices)

spare/utils.py

@@ -2,6 +2,21 @@
 import math
 import numpy as np
 
+
+
+def maybe_init(class_or_insatnce):
+ if isinstance(class_or_insatnce, type):
+ return class_or_insatnce()
+ else:
+ return class_or_insatnce
+
+def load_backend(backend):
+ if backend=="torch":
+ from spare.backend_torch import TorchBackend
+ return TorchBackend()
+ else:
+ RuntimeError("Only torch backend is currently supported")
+
 def get_best_np_dtype_for(min_value, max_value):
  uint32_bounds = np.iinfo("uint32")
  uint64_bounds = np.iinfo("uint64")

spare/weighting_model.py

@@ -61,7 +61,7 @@ def _get_vars_to_save(self):
  return super()._get_vars_to_save() | {
  "k1": self.k1,
  "b": self.b,
- "idf_weighting": idf_weighting
+ "idf_weighting": self.idf_weighting
  } 
 
  def _load_vars(self, data):