From a0e0c71ea42a71efa1965fd4c6a7382a1d2ad99b Mon Sep 17 00:00:00 2001 From: Khoroshevskyi Date: Wed, 17 Dec 2025 08:49:35 -0500 Subject: [PATCH 01/11] Updating model initialization and qdrant connection --- .github/workflows/python-publish.yml | 2 +- bbconf/_version.py | 2 +- bbconf/config_parser/bedbaseconfig.py | 369 +++++++++++++++++--------- bbconf/config_parser/const.py | 9 +- bbconf/config_parser/models.py | 14 +- bbconf/modules/bedfiles.py | 284 +++++++++++++++++--- manual_testing.py | 7 +- requirements/requirements-all.txt | 2 +- tests/config_test.yaml | 2 +- 9 files changed, 509 insertions(+), 182 deletions(-) diff --git a/.github/workflows/python-publish.yml b/.github/workflows/python-publish.yml index e1da3429..2516e668 100644 --- a/.github/workflows/python-publish.yml +++ b/.github/workflows/python-publish.yml @@ -13,7 +13,7 @@ jobs: id-token: write # IMPORTANT: this permission is mandatory for trusted publishing steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v5 - name: Set up Python uses: actions/setup-python@v5 with: diff --git a/bbconf/_version.py b/bbconf/_version.py index f23a6b39..9e78220f 100644 --- a/bbconf/_version.py +++ b/bbconf/_version.py @@ -1 +1 @@ -__version__ = "0.13.0" +__version__ = "0.14.0" diff --git a/bbconf/config_parser/bedbaseconfig.py b/bbconf/config_parser/bedbaseconfig.py index d3244b33..87a76a41 100644 --- a/bbconf/config_parser/bedbaseconfig.py +++ b/bbconf/config_parser/bedbaseconfig.py @@ -3,7 +3,6 @@ import warnings from pathlib import Path from typing import List, Literal, Union -import numpy as np import boto3 import qdrant_client @@ -11,6 +10,11 @@ import s3fs import yacman import zarr +import requests + +from fastembed import TextEmbedding +from sentence_transformers import SparseEncoder + from botocore.exceptions import BotoCoreError, EndpointConnectionError from geniml.region2vec.main import Region2VecExModel from geniml.search import BED2BEDSearchInterface @@ -28,7 +32,6 @@ S3_BEDSET_PATH_FOLDER, S3_FILE_PATH_FOLDER, S3_PLOTS_PATH_FOLDER, - TEXT_EMBEDDING_DIMENSION, ) from bbconf.config_parser.models import ConfigFile from bbconf.const import PKG_NAME, ZARR_TOKENIZED_FOLDER @@ -64,24 +67,54 @@ def __init__(self, config: Union[Path, str], init_ml: bool = True): self._config = self._read_config_file(self.cfg_path) self._db_engine = self._init_db_engine() - self._qdrant_engine = self._init_qdrant_backend() - self._qdrant_text_engine = self._init_qdrant_text_backend() - self._qdrant_advanced_engine = self._init_qdrant_advanced_backend() + try: + self.qdrant_client: QdrantClient = self._init_qdrant_client() + except Exception as err: + _LOGGER.error( + f"Unable to create Qdrant client. Skipping ML model initialization. Error: {err}" + ) + init_ml = False if init_ml: - self._b2bsi = self._init_b2bsi_object() - self._r2v = self._init_r2v_object() - self._bivec = self._init_bivec_object() - self._umap_model: Union[UMAP, None] = self._init_umap_model() + + self.dense_encoder: TextEmbedding = self._init_dense_encoder() + self.sparce_encoder: Union[SparseEncoder, None] = self._init_sparce_model() + self._umap_encoder: Union[UMAP, None] = self._init_umap_model() + self.r2v_encoder = self._init_r2v_encoder() + + self._init_qdrant_hybrid( + qdrant_cl=self.qdrant_client, + dense_encoder=self.dense_encoder, + ) + + self.qdrant_file_backend: Union[QdrantBackend, None] = ( + self._init_qdrant_file_backend(qdrant_cl=self.qdrant_client) + ) # used for bivec search + self._qdrant_text_backend: Union[QdrantBackend, None] = ( + self._init_qdrant_text_backend( + qdrant_cl=self.qdrant_client, + dense_encoder=self.dense_encoder, + ) + ) # used for bivec search + + self._b2b_search_interface = self._init_b2b_search_interface( + qdrant_file_backend=self.qdrant_file_backend, + ) + + self._bivec_search_interface = self._init_bivec_interface( + qdrant_file_backend=self.qdrant_file_backend, + qdrant_text_backend=self._qdrant_text_backend, + text_encoder=self.dense_encoder, + ) else: _LOGGER.info( "Skipping initialization of ML models, init_ml parameter set to False." ) - - self._b2bsi = None - self._r2v = None - self._bivec = None - self._umap_model: Union[UMAP, None] = None + self.r2v_encoder = None + self._b2b_search_interface = None + self._bivec_search_interface = None + self._umap_encoder: Union[UMAP, None] = None + self.sparce_encoder = None self._phc = self._init_pephubclient() self._boto3_client = self._init_boto3_client() @@ -135,16 +168,7 @@ def b2bsi(self) -> Union[BED2BEDSearchInterface, None]: :return: bed2bednn object """ - return self._b2bsi - - @property - def r2v(self) -> Region2VecExModel: - """ - Get region2vec object - - :return: region2vec object - """ - return self._r2v + return self._b2b_search_interface @property def bivec(self) -> BiVectorSearchInterface: @@ -154,16 +178,16 @@ def bivec(self) -> BiVectorSearchInterface: :return: bivec search interface object """ - return self._bivec + return self._bivec_search_interface - @property - def qdrant_engine(self) -> QdrantBackend: - """ - Get qdrant engine - - :return: qdrant engine - """ - return self._qdrant_engine + # @property + # def qdrant_engine(self) -> QdrantBackend: + # """ + # Get qdrant engine + # + # :return: qdrant engine + # """ + # return self._qdrant_file_backend @property def phc(self) -> PEPHubClient: @@ -226,81 +250,136 @@ def _init_db_engine(self) -> BaseEngine: drivername=f"{self._config.database.dialect}+{self._config.database.driver}", ) - def _init_qdrant_backend(self) -> QdrantBackend: + def _init_qdrant_client(self) -> QdrantClient: """ Create qdrant client object using credentials provided in config file - - :return: QdrantClient """ - _LOGGER.info("Initializing qdrant engine...") + _LOGGER.info("Initializing qdrant client...") + try: - return QdrantBackend( - collection=self._config.qdrant.file_collection, - qdrant_host=self._config.qdrant.host, - qdrant_port=self._config.qdrant.port, - qdrant_api_key=self._config.qdrant.api_key, + qdrant_cl = QdrantClient( + url=self.config.qdrant.host, + port=self.config.qdrant.port, + api_key=self.config.qdrant.api_key, ) + except qdrant_client.http.exceptions.ResponseHandlingException as err: + raise BedBaseConfError( + f"Error in Connection to qdrant! skipping... Error: {err}" + ) + + return qdrant_cl + + def _init_qdrant_file_backend( + self, qdrant_cl: QdrantClient + ) -> Union[QdrantBackend, None]: + """ + Create qdrant client object using credentials provided in config file + + :param: qdrant_cl: QdrantClient object + :return: QdrantClient + """ + + _LOGGER.info("Initializing qdrant bivec file backend...") + + if not isinstance(qdrant_cl, QdrantClient): _LOGGER.error( - f"Error in Connection to qdrant! skipping... Error: {err}. Qdrant host: {self._config.qdrant.host}" + f"Unable to create Qdrant bivec file collection, qdrant client is None." ) - warnings.warn( - f"error in Connection to qdrant! skipping... Error: {err}", UserWarning + return None + + try: + return QdrantBackend( + qdrant_client=qdrant_cl, + collection=self.config.qdrant.file_collection, ) + except Exception as e: + _LOGGER.error(f"Unable to create Qdrant collection: {e}") + return None - def _init_qdrant_text_backend(self) -> Union[QdrantBackend, None]: + def _init_qdrant_text_backend( + self, qdrant_cl: QdrantClient, dense_encoder: TextEmbedding + ) -> Union[QdrantBackend, None]: """ Create qdrant client text embedding object using credentials provided in config file + :param: qdrant_cl: QdrantClient object + :param: dense_encoder: TextEmbedding model for encoding text queries :return: QdrantClient """ - _LOGGER.info("Initializing qdrant text engine...") + _LOGGER.info("Initializing qdrant bivec text backend...") + + if not isinstance(qdrant_cl, QdrantClient): + _LOGGER.error( + f"Unable to create Qdrant bivec text collection, qdrant client is None." + ) + return None + if not isinstance(dense_encoder, TextEmbedding): + _LOGGER.error( + f"Unable to create Qdrant bivec text collection,, dense encoder is None." + ) + return None + + dimensions = int(dense_encoder.get_embedding_size(self._config.path.text2vec)) try: return QdrantBackend( - dim=TEXT_EMBEDDING_DIMENSION, + qdrant_client=qdrant_cl, + dim=dimensions, collection=self.config.qdrant.text_collection, - qdrant_host=self.config.qdrant.host, - qdrant_api_key=self.config.qdrant.api_key, ) except Exception as e: - _LOGGER.error( - f"Error in Connection to qdrant text! skipping {e}. Qdrant host: {self._config.qdrant.host}" - ) - warnings.warn( - "Error in Connection to qdrant text! skipping...", UserWarning - ) + _LOGGER.error(f"Unable to create Qdrant collection: {e}") return None - def _init_qdrant_advanced_backend(self) -> Union[QdrantClient, None]: + def _init_qdrant_hybrid( + self, qdrant_cl: QdrantClient, dense_encoder: TextEmbedding + ) -> None: """ - Create qdrant client text embedding object using credentials provided in config file + Create qdrant client with sparse and text embedding object using credentials provided in config file + :param: qdrant_cl: QdrantClient object + :param: dense_encoder: TextEmbedding model for encoding text queries :return: QdrantClient """ - COLLECTION_NAME = self.config.qdrant.search_collection - DIMENSIONS = 384 + _LOGGER.info("Initializing qdrant sparse collection...") - _LOGGER.info("Initializing qdrant text advanced engine...") - - try: - qdrant_cl = QdrantClient( - url=self.config.qdrant.host, - port=self.config.qdrant.port, - api_key=self.config.qdrant.api_key, + if not isinstance(qdrant_cl, QdrantClient): + _LOGGER.error( + f"Unable to create Qdrant hybrid collection, qdrant client is None." ) + return None + if not isinstance(dense_encoder, TextEmbedding): + _LOGGER.error( + f"Unable to create Qdrant hybrid collection, dense encoder is None." + ) + return None - if not qdrant_cl.collection_exists(COLLECTION_NAME): + dimensions = int(dense_encoder.get_embedding_size(self._config.path.text2vec)) + collection_name = self.config.qdrant.search_collection + + try: + if not qdrant_cl.collection_exists(collection_name): _LOGGER.info( "Collection 'bedbase_query_search' does not exist, creating it." ) + qdrant_cl.create_collection( - collection_name=COLLECTION_NAME, - vectors_config=models.VectorParams( - size=DIMENSIONS, distance=models.Distance.COSINE - ), + collection_name=collection_name, + vectors_config={ + "dense": models.VectorParams( + size=dimensions, distance=models.Distance.COSINE + ), + }, + sparse_vectors_config={ + "sparse": models.SparseVectorParams( + index=models.SparseIndexParams( + on_disk=False, + ) + ) + }, quantization_config=models.ScalarQuantization( scalar=models.ScalarQuantizationConfig( type=models.ScalarType.INT8, @@ -311,55 +390,64 @@ def _init_qdrant_advanced_backend(self) -> Union[QdrantClient, None]: ) qdrant_cl.create_payload_index( - collection_name=COLLECTION_NAME, + collection_name=collection_name, field_name="assay", - field_schema="keyword", + field_type=models.PayloadSchemaType.KEYWORD, ) qdrant_cl.create_payload_index( - collection_name=COLLECTION_NAME, + collection_name=collection_name, field_name="genome_alias", - field_schema="keyword", + field_type=models.PayloadSchemaType.KEYWORD, ) - return qdrant_cl - - except qdrant_client.http.exceptions.ResponseHandlingException as err: + except Exception as err: _LOGGER.error( - f"Error in Connection to qdrant! skipping... Error: {err}. Qdrant host: {self._config.qdrant.host}" + f"Error in creating Qdrant hybrid collection! skipping... Error: {err}. Qdrant host: {self._config.qdrant.host}" ) warnings.warn( - f"error in Connection to qdrant! skipping... Error: {err}", UserWarning + f"error in creating Qdrant hybrid collection! skipping... Error: {err}", + UserWarning, ) return None - def _init_bivec_object(self) -> Union[BiVectorSearchInterface, None]: + def _init_bivec_interface( + self, + qdrant_file_backend: QdrantBackend, + qdrant_text_backend: QdrantBackend, + text_encoder: TextEmbedding, + ) -> Union[BiVectorSearchInterface, None]: """ Create BiVectorSearchInterface object using credentials provided in config file + :param: qdrant_file_backend: QdrantBackend for file vectors + :param: qdrant_text_backend: QdrantBackend for text vectors + :param: text_encoder: TextEmbedding model for encoding text queries :return: BiVectorSearchInterface """ _LOGGER.info("Initializing BiVectorBackend...") search_backend = BiVectorBackend( - metadata_backend=self._qdrant_text_engine, bed_backend=self._qdrant_engine + metadata_backend=qdrant_text_backend, bed_backend=qdrant_file_backend ) _LOGGER.info("Initializing BiVectorSearchInterface...") search_interface = BiVectorSearchInterface( backend=search_backend, - query2vec=self.config.path.text2vec, + query2vec=text_encoder, ) return search_interface - def _init_b2bsi_object(self) -> Union[BED2BEDSearchInterface, None]: + def _init_b2b_search_interface( + self, qdrant_file_backend: QdrantBackend + ) -> Union[BED2BEDSearchInterface, None]: """ Create Bed 2 BED search interface and return this object :return: Bed2BEDSearchInterface object """ try: - _LOGGER.info("Initializing search interfaces...") + _LOGGER.info("Initializing search bed 2 bed search interfaces...") return BED2BEDSearchInterface( - backend=self.qdrant_engine, + backend=qdrant_file_backend, query2vec=BED2Vec(model=self._config.path.region2vec), ) except Exception as e: @@ -370,56 +458,47 @@ def _init_b2bsi_object(self) -> Union[BED2BEDSearchInterface, None]: ) return None - @staticmethod - def _init_pephubclient() -> Union[PEPHubClient, None]: - """ - Create Pephub client object using credentials provided in config file - - :return: PephubClient - """ - - # try: - # _LOGGER.info("Initializing PEPHub client...") - # return PEPHubClient() - # except Exception as e: - # _LOGGER.error(f"Error in creating PephubClient object: {e}") - # warnings.warn(f"Error in creating PephubClient object: {e}", UserWarning) - # return None - return None - - def _init_boto3_client( - self, - ) -> boto3.client: + def _init_r2v_encoder(self) -> Union[Region2VecExModel, None]: """ - Create Pephub client object using credentials provided in config file - - :return: PephubClient + Create Region2VecExModel object using credentials provided in config file """ try: - return boto3.client( - "s3", - endpoint_url=self._config.s3.endpoint_url, - aws_access_key_id=self._config.s3.aws_access_key_id, - aws_secret_access_key=self._config.s3.aws_secret_access_key, + _LOGGER.info( + f"Initializing region2vec encoder... Model used: {self.config.path.region2vec}" ) + return Region2VecExModel(self.config.path.region2vec) except Exception as e: - _LOGGER.error(f"Error in creating boto3 client object: {e}") - warnings.warn(f"Error in creating boto3 client object: {e}", UserWarning) + _LOGGER.error(f"Error in creating Region2VecExModel object: {e}") + warnings.warn( + f"Error in creating Region2VecExModel object: {e}", UserWarning + ) return None - def _init_r2v_object(self) -> Union[Region2VecExModel, None]: + def _init_dense_encoder(self) -> Union[None, TextEmbedding]: """ - Create Region2VecExModel object using credentials provided in config file + Initialize dense model from the specified path or huggingface model hub + """ + + _LOGGER.info( + f"Initializing dense encoder... Model used: {self.config.path.text2vec}" + ) + dense_encoder = TextEmbedding(self.config.path.text2vec) + return dense_encoder + + def _init_sparce_model(self) -> Union[None, SparseEncoder]: + """ + Initialize SparseEncoder model from the specified path or huggingface model hub """ try: - _LOGGER.info("Initializing R2V object...") - return Region2VecExModel(self.config.path.region2vec) - except Exception as e: - _LOGGER.error(f"Error in creating Region2VecExModel object: {e}") - warnings.warn( - f"Error in creating Region2VecExModel object: {e}", UserWarning + _LOGGER.info( + f"Initializing sparse encoder... Model used: {self.config.path.sparse_model}" ) + sparse_encoder = SparseEncoder(self.config.path.sparse_model) + except Exception as e: + _LOGGER.error(f"Error in creating SparseEncoder object: {e}") + warnings.warn(f"Error in creating SparseEncoder object: {e}", UserWarning) return None + return sparse_encoder def _init_umap_model(self) -> Union[UMAP, None]: """ @@ -433,9 +512,8 @@ def _init_umap_model(self) -> Union[UMAP, None]: return None model_path = self.config.path.umap_model - + umap_model = None if model_path.startswith(("http://", "https://")): - import requests try: response = requests.get(model_path) @@ -462,6 +540,26 @@ def _init_umap_model(self) -> Union[UMAP, None]: umap_model.random_state = 42 return umap_model + def _init_boto3_client( + self, + ) -> Union[boto3.client, None]: + """ + Create Pephub client object using credentials provided in config file + + :return: PephubClient + """ + try: + return boto3.client( + "s3", + endpoint_url=self._config.s3.endpoint_url, + aws_access_key_id=self._config.s3.aws_access_key_id, + aws_secret_access_key=self._config.s3.aws_secret_access_key, + ) + except Exception as e: + _LOGGER.error(f"Error in creating boto3 client object: {e}") + warnings.warn(f"Error in creating boto3 client object: {e}", UserWarning) + return None + def upload_s3(self, file_path: str, s3_path: Union[Path, str]) -> None: """ Upload file to s3. @@ -578,6 +676,23 @@ def delete_files_s3(self, files: List[FileModel]) -> None: self.delete_s3(file.path_thumbnail) return None + @staticmethod + def _init_pephubclient() -> Union[PEPHubClient, None]: + """ + Create Pephub client object using credentials provided in config file + + :return: PephubClient + """ + + # try: + # _LOGGER.info("Initializing PEPHub client...") + # return PEPHubClient() + # except Exception as e: + # _LOGGER.error(f"Error in creating PephubClient object: {e}") + # warnings.warn(f"Error in creating PephubClient object: {e}", UserWarning) + # return None + return None + def get_prefixed_uri(self, postfix: str, access_id: str) -> str: """ Return uri with correct prefix (schema) diff --git a/bbconf/config_parser/const.py b/bbconf/config_parser/const.py index b9541f51..61aad4ec 100644 --- a/bbconf/config_parser/const.py +++ b/bbconf/config_parser/const.py @@ -5,16 +5,17 @@ DEFAULT_QDRANT_HOST = "localhost" DEFAULT_QDRANT_PORT = 6333 -DEFAULT_QDRANT_COLLECTION_NAME = "bedbase" -DEFAULT_QDRANT_TEXT_COLLECTION_NAME = "bed_text" -DEFAULT_QDRANT_SEARCH_COLLECTION_NAME = "bedbase_query_search" +DEFAULT_QDRANT_FILE_COLLECTION_NAME = "bedbase" +DEFAULT_QDRANT_BIVEC_COLLECTION_NAME = "bed_text" +DEFAULT_QDRANT_HYBRID_COLLECTION_NAME = "bedbase_query_search" DEFAULT_QDRANT_API_KEY = None DEFAULT_SERVER_PORT = 80 DEFAULT_SERVER_HOST = "0.0.0.0" DEFAULT_TEXT2VEC_MODEL = "sentence-transformers/all-MiniLM-L6-v2" -DEFAULT_REGION2_VEC_MODEL = "databio/r2v-ChIP-atlas-hg38" +DEFAULT_SPARSE_MODEL = "prithivida/Splade_PP_en_v2" +DEFAULT_REGION2_VEC_MODEL = "databio/r2v_encoder-ChIP-atlas-hg38" DEFAULT_PEPHUB_NAMESPACE = "databio" DEFAULT_PEPHUB_NAME = "bedbase_all" diff --git a/bbconf/config_parser/models.py b/bbconf/config_parser/models.py index 2e0c14cc..fcfaedbd 100644 --- a/bbconf/config_parser/models.py +++ b/bbconf/config_parser/models.py @@ -13,15 +13,16 @@ DEFAULT_PEPHUB_NAME, DEFAULT_PEPHUB_NAMESPACE, DEFAULT_PEPHUB_TAG, - DEFAULT_QDRANT_COLLECTION_NAME, + DEFAULT_QDRANT_FILE_COLLECTION_NAME, DEFAULT_QDRANT_PORT, - DEFAULT_QDRANT_TEXT_COLLECTION_NAME, + DEFAULT_QDRANT_BIVEC_COLLECTION_NAME, DEFAULT_REGION2_VEC_MODEL, DEFAULT_S3_BUCKET, DEFAULT_SERVER_HOST, DEFAULT_SERVER_PORT, DEFAULT_TEXT2VEC_MODEL, - DEFAULT_QDRANT_SEARCH_COLLECTION_NAME, + DEFAULT_SPARSE_MODEL, + DEFAULT_QDRANT_HYBRID_COLLECTION_NAME, ) _LOGGER = logging.getLogger(__name__) @@ -53,9 +54,9 @@ class ConfigQdrant(BaseModel): host: str port: int = DEFAULT_QDRANT_PORT api_key: Optional[str] = None - file_collection: str = DEFAULT_QDRANT_COLLECTION_NAME - text_collection: Optional[str] = DEFAULT_QDRANT_TEXT_COLLECTION_NAME - search_collection: Optional[str] = DEFAULT_QDRANT_SEARCH_COLLECTION_NAME + file_collection: str = DEFAULT_QDRANT_FILE_COLLECTION_NAME + text_collection: Optional[str] = DEFAULT_QDRANT_BIVEC_COLLECTION_NAME + search_collection: Optional[str] = DEFAULT_QDRANT_HYBRID_COLLECTION_NAME class ConfigServer(BaseModel): @@ -67,6 +68,7 @@ class ConfigPath(BaseModel): region2vec: str = DEFAULT_REGION2_VEC_MODEL # vec2vec: str = DEFAULT_VEC2VEC_MODEL text2vec: str = DEFAULT_TEXT2VEC_MODEL + sparse_model: str = DEFAULT_SPARSE_MODEL umap_model: Union[str, None] = None # Path or link to pre-trained UMAP model diff --git a/bbconf/modules/bedfiles.py b/bbconf/modules/bedfiles.py index 70633c68..cf83775c 100644 --- a/bbconf/modules/bedfiles.py +++ b/bbconf/modules/bedfiles.py @@ -19,8 +19,8 @@ from sqlalchemy.orm.attributes import flag_modified from sqlalchemy.dialects import postgresql from tqdm import tqdm -from fastembed import TextEmbedding from qdrant_client import models +from qdrant_client.models import QueryResponse from bbconf.config_parser.bedbaseconfig import BedBaseConfig from bbconf.const import DEFAULT_LICENSE, PKG_NAME, ZARR_TOKENIZED_FOLDER @@ -85,14 +85,10 @@ def __init__(self, config: BedBaseConfig, bbagent_obj=None): """ self._sa_engine = config.db_engine.engine self._db_engine = config.db_engine - self._qdrant_engine = config.qdrant_engine self._boto3_client = config.boto3_client self._config = config self.bb_agent = bbagent_obj - self._embedding_model = TextEmbedding(config.config.path.text2vec) - # self._embedding_model = TextEmbedding("BAAI/bge-large-en-v1.5") - def get(self, identifier: str, full: bool = False) -> BedMetadataAll: """ Get file metadata by identifier. @@ -272,7 +268,7 @@ def get_neighbours( if not self.exists(identifier): raise BEDFileNotFoundError(f"Bed file with id: {identifier} not found.") s = identifier - results = self._qdrant_engine.qd_client.query_points( + results = self._config.qdrant_file_backend.qd_client.query_points( collection_name=self._config.config.qdrant.file_collection, query="-".join([s[:8], s[8:12], s[12:16], s[16:20], s[20:]]), limit=limit, @@ -390,7 +386,7 @@ def get_embedding(self, identifier: str) -> BedEmbeddingResult: """ if not self.exists(identifier): raise BEDFileNotFoundError(f"Bed file with id: {identifier} not found.") - result = self._qdrant_engine.qd_client.retrieve( + result = self._config.qdrant_file_backend.qd_client.retrieve( collection_name=self._config.config.qdrant.file_collection, ids=[identifier], with_vectors=True, @@ -1172,12 +1168,12 @@ def upload_file_qdrant( _LOGGER.debug(f"Adding bed file to qdrant. bed_id: {bed_id}") - if not isinstance(self._qdrant_engine, QdrantBackend): + if not isinstance(self._config.qdrant_file_backend, QdrantBackend): raise QdrantInstanceNotInitializedError("Could not upload file.") bed_embedding = self._embed_file(bed_file) - self._qdrant_engine.load( + self._config.qdrant_file_backend.load( ids=[bed_id], vectors=bed_embedding, payloads=[{**payload}], @@ -1193,9 +1189,9 @@ def _embed_file(self, bed_file: Union[str, GRegionSet]) -> np.ndarray: :return np array of embeddings """ - if self._qdrant_engine is None: + if self._config.qdrant_file_backend is None: raise QdrantInstanceNotInitializedError - if not self._config.r2v: + if not self._config.r2v_encoder: raise BedBaseConfError( "Could not add add region to qdrant. Invalid type, or path. " ) @@ -1212,7 +1208,7 @@ def _embed_file(self, bed_file: Union[str, GRegionSet]) -> np.ndarray: raise BedBaseConfError( "Could not add add region to qdrant. Invalid type, or path. " ) - bed_embedding = np.mean(self._config.r2v.encode(bed_region_set), axis=0) + bed_embedding = np.mean(self._config.r2v_encoder.encode(bed_region_set), axis=0) vec_dim = bed_embedding.shape[0] return bed_embedding.reshape(1, vec_dim) @@ -1223,11 +1219,11 @@ def _get_umap_file(self, bed_file: Union[str, GRegionSet]) -> np.ndarray: :param bed_file: bed file path or region set """ - if self._config._umap_model is None: + if self._config._umap_encoder is None: raise BedBaseConfError("UMAP model is not initialized.") bed_embedding = self._embed_file(bed_file) - bed_umap = self._config._umap_model.transform(bed_embedding) + bed_umap = self._config._umap_encoder.transform(bed_embedding) return bed_umap def text_to_bed_search( @@ -1271,7 +1267,7 @@ def text_to_bed_search( ) if with_metadata: - count = self._config._qdrant_advanced_engine.get_collection( + count = self._config.qdrant_client.get_collection( collection_name=self._config.config.qdrant.file_collection ).points_count else: @@ -1501,7 +1497,7 @@ def reindex_qdrant(self, batch: int = 100, purge: bool = False) -> None: pbar.set_description( "Uploading points to qdrant using batch..." ) - operation_info = self._config.qdrant_engine.qd_client.upsert( + operation_info = self._config.qdrant_file_backend.qd_client.upsert( collection_name=self._config.config.qdrant.file_collection, points=points_list, ) @@ -1515,7 +1511,7 @@ def reindex_qdrant(self, batch: int = 100, purge: bool = False) -> None: pbar.update(1) _LOGGER.info("Uploading points to qdrant using batches...") - operation_info = self._config.qdrant_engine.qd_client.upsert( + operation_info = self._config.qdrant_file_backend.qd_client.upsert( collection_name=self._config.config.qdrant.file_collection, points=points_list, ) @@ -1530,22 +1526,13 @@ def delete_qdrant_point(self, identifier: str) -> None: :return: None """ - result = self._config.qdrant_engine.qd_client.delete( + result = self._config.qdrant_file_backend.qd_client.delete( collection_name=self._config.config.qdrant.file_collection, points_selector=PointIdsList( points=[identifier], ), ) - return result - - def create_qdrant_collection(self) -> bool: - """ - Create qdrant collection for bed files. - """ - return self._config.qdrant_engine.qd_client.create_collection( - collection_name=self._config.config.qdrant.file_collection, - vectors_config=VectorParams(size=100, distance=Distance.DOT), - ) + return None def exists(self, identifier: str) -> bool: """ @@ -2006,7 +1993,103 @@ def _update_sources( session.commit() - def reindex_semantic_search(self, batch: int = 1000, purge: bool = False) -> None: + # def reindex_semantic_search(self, batch: int = 1000, purge: bool = False) -> None: + # """ + # Reindex all bed files for semantic database + # + # :param batch: number of files to upload in one batch + # :param purge: resets indexed in database for all files to False + # + # :return: None + # """ + # + # # Add column that will indicate if this file is indexed or not + # statement = ( + # select(Bed) + # .join(BedMetadata, Bed.id == BedMetadata.id) + # .where(Bed.indexed == False) + # .limit(150000) + # ) + # + # with Session(self._sa_engine) as session: + # + # if purge: + # _LOGGER.info("Purging indexed files in the database ...") + # session.query(Bed).update({Bed.indexed: False}) + # session.commit() + # _LOGGER.info("Purged indexed files in the database successfully!") + # + # _LOGGER.info("Fetching data from the database ...") + # results = session.scalars(statement) + # + # _LOGGER.info("Fetch data successfully!") + # + # points = [] + # results = [result for result in results] + # + # with tqdm(total=len(results), position=0, leave=True) as pbar: + # processed_number = 0 + # for result in results: + # text = ( + # f"biosample is {result.annotations.cell_line} / {result.annotations.cell_type} / " + # f"{result.annotations.tissue} with target {result.annotations.target} " + # f"assay {result.annotations.assay}." + # f"File name {result.name} with summary {result.description}" + # ) + # + # embeddings_list = list(self._config.dense_encoder.embed(text)) + # # result_list.append( + # data = VectorMetadata( + # id=result.id, + # name=result.name, + # description=result.description, + # genome_alias=result.genome_alias, + # genome_digest=result.genome_digest, + # cell_line=result.annotations.cell_line, + # cell_type=result.annotations.cell_type, + # tissue=result.annotations.tissue, + # target=result.annotations.target, + # treatment=result.annotations.treatment, + # assay=result.annotations.assay, + # species_name=result.annotations.species_name, + # ) + # + # points.append( + # PointStruct( + # id=result.id, + # vector=list(embeddings_list[0]), + # payload=data.model_dump(), + # ) + # ) + # processed_number += 1 + # result.indexed = True + # + # if processed_number % batch == 0: + # pbar.set_description( + # "Uploading points to qdrant using batch..." + # ) + # operation_info = self._config._qdrant_advanced_engine.upsert( + # collection_name=self._config.config.qdrant.search_collection, + # points=points, + # ) + # session.commit() + # pbar.write("Uploaded batch to qdrant.") + # points = [] + # assert operation_info.status == "completed" + # + # pbar.write(f"File: {result.id} successfully indexed.") + # pbar.update(1) + # + # operation_info = self._config._qdrant_advanced_engine.upsert( + # collection_name=self._config.config.qdrant.search_collection, + # points=points, + # ) + # assert operation_info.status == "completed" + # session.commit() + # + # return None + + def reindex_hybrid_search(self, batch: int = 1000, purge: bool = False) -> None: """ Reindex all bed files for semantic database @@ -2021,7 +2104,7 @@ def reindex_semantic_search(self, batch: int = 1000, purge: bool = False) -> Non select(Bed) .join(BedMetadata, Bed.id == BedMetadata.id) .where(Bed.indexed == False) - .limit(150000) + .limit(150) ) with Session(self._sa_engine) as session: @@ -2050,7 +2133,26 @@ def reindex_semantic_search(self, batch: int = 1000, purge: bool = False) -> Non f"File name {result.name} with summary {result.description}" ) - embeddings_list = list(self._embedding_model.embed(text)) + embeddings_list = list(self._config.dense_encoder.embed(text)) + + if self._config.sparce_encoder: + sparse_result = self._config.sparce_encoder.encode( + text + ).coalesce() + + sparse_embeddings = models.SparseVector( + indices=sparse_result.indices().tolist()[0], + values=sparse_result.values().tolist(), + ) + + point_vectors = { + "dense": list(embeddings_list[0]), + "sparse": sparse_embeddings, + } + else: + point_vectors = { + "dense": list(embeddings_list[0]), + } # result_list.append( data = VectorMetadata( id=result.id, @@ -2070,7 +2172,7 @@ def reindex_semantic_search(self, batch: int = 1000, purge: bool = False) -> Non points.append( PointStruct( id=result.id, - vector=list(embeddings_list[0]), + vector=point_vectors, payload=data.model_dump(), ) ) @@ -2081,7 +2183,7 @@ def reindex_semantic_search(self, batch: int = 1000, purge: bool = False) -> Non pbar.set_description( "Uploading points to qdrant using batch..." ) - operation_info = self._config._qdrant_advanced_engine.upsert( + operation_info = self._config.qdrant_client.upsert( collection_name=self._config.config.qdrant.search_collection, points=points, ) @@ -2093,7 +2195,7 @@ def reindex_semantic_search(self, batch: int = 1000, purge: bool = False) -> Non pbar.write(f"File: {result.id} successfully indexed.") pbar.update(1) - operation_info = self._config._qdrant_advanced_engine.upsert( + operation_info = self._config.qdrant_client.upsert( collection_name=self._config.config.qdrant.search_collection, points=points, ) @@ -2142,11 +2244,11 @@ def semantic_search( ) ) - embeddings_list = list(self._embedding_model.embed(query))[0] + embeddings_list = list(self._config.dense_encoder.embed(query))[0] - results = self._config._qdrant_advanced_engine.search( + results: QueryResponse = self._config.qdrant_client.query_points( collection_name=self._config.config.qdrant.search_collection, - query_vector=list(embeddings_list), + query=list(embeddings_list), limit=limit, offset=offset, search_params=models.SearchParams( @@ -2161,7 +2263,113 @@ def semantic_search( ) result_list = [] - for result in results: + for result in results.points: + result_id = result.id.replace("-", "") + + if with_metadata: + metadata = self.get(result_id, full=False) + else: + metadata = None + + result_list.append( + QdrantSearchResult( + id=result_id, + payload=result.payload, + score=result.score, + metadata=metadata, + ) + ) + + if with_metadata: + count = self.bb_agent.get_stats().bedfiles_number + else: + count = 0 + + return BedListSearchResult( + count=count, + limit=limit, + offset=offset, + results=result_list, + ) + + def hybrid_search( + self, + query: str = "liver", + genome_alias: str = "", + assay: str = "", + limit: int = 100, + offset: int = 0, + with_metadata: bool = True, + ) -> BedListSearchResult: + """ + Run semantic search for bed files using qdrant. + This is not bivec search, but usual qdrant search with sparse and dense embeddings. + + :param query: text query to search for + :param genome_alias: genome alias to filter results + :param assay: filter by assay type + :param limit: number of results to return + :param offset: offset to start from + :param with_metadata: if True, metadata will be returned in the results. Default is True. + + :return: list of bed file metadata + """ + + must_statement = [] + + if genome_alias: + must_statement.append( + models.FieldCondition( + key="genome_alias", + match=models.MatchValue(value=genome_alias), + ) + ) + if assay: + must_statement.append( + models.FieldCondition( + key="assay", + match=models.MatchValue(value=assay), + ) + ) + + dense_query = list(list(self._config.dense_encoder.embed(query))[0]) + if self._config.sparce_encoder: + sparse_result = self._config.sparce_encoder.encode(query).coalesce() + sparse_embeddings = models.SparseVector( + indices=sparse_result.indices().tolist()[0], + values=sparse_result.values().tolist(), + ) + + hybrid_query = [ + # Dense retrieval: semantic understanding + models.Prefetch(query=dense_query, using="dense", limit=limit), + # Sparse retrieval: exact technical term matching + models.Prefetch(query=sparse_embeddings, using="sparse", limit=limit), + ] + else: + hybrid_query = [ + # Dense retrieval: semantic understanding + models.Prefetch(query=dense_query, using="dense", limit=limit), + ] + + results = self._config.qdrant_client.query_points( + collection_name=self._config.config.qdrant.search_collection, + limit=limit, + offset=offset, + prefetch=hybrid_query, + query=models.FusionQuery(fusion=models.Fusion.RRF), + with_payload=True, + with_vectors=True, + search_params=models.SearchParams( + exact=True, + ), + query_filter=( + models.Filter(must=must_statement) if must_statement else None + ), + ) + + result_list = [] + for result in results.points: result_id = result.id.replace("-", "") if with_metadata: diff --git a/manual_testing.py b/manual_testing.py index 2fc2205d..0ecae33a 100644 --- a/manual_testing.py +++ b/manual_testing.py @@ -329,8 +329,9 @@ def new_search(): agent = BedBaseAgent(config="/home/bnt4me/virginia/repos/bedhost/config.yaml") time1 = time.time() - results = agent.bed.reindex_semantic_search() + # results = agent.bed.reindex_semantic_search() # results = agent.bed.comp_search() + results = agent.bed.hybrid_search("leukemia") time2 = time.time() print(f"Time taken: {time2 - time1} seconds") @@ -370,10 +371,10 @@ def reindex_files(): # neighbour_beds() # sql_search() # config_t() - compreh_stats() + # compreh_stats() # get_unprocessed_files() # get_genomes() - # new_search() + new_search() # external_search() # get_assay_list() diff --git a/requirements/requirements-all.txt b/requirements/requirements-all.txt index f7fc476d..cbf01e83 100644 --- a/requirements/requirements-all.txt +++ b/requirements/requirements-all.txt @@ -1,7 +1,7 @@ yacman >= 0.9.1 sqlalchemy >= 2.0.0 gtars >= 0.4.0 -geniml[ml] >= 0.7.1 +geniml[ml] >= 0.8.2 psycopg >= 3.1.15 coloredlogs pydantic >= 2.9.0 diff --git a/tests/config_test.yaml b/tests/config_test.yaml index 54c3ba17..ef8069eb 100644 --- a/tests/config_test.yaml +++ b/tests/config_test.yaml @@ -1,5 +1,5 @@ path: - region2vec: databio/r2v-encode-hg38 + region2vec: databio/r2v_encoder-encode-hg38 # vec2vec: databio/v2v-geo-hg38 database: host: localhost From b36de64b077a6a82aa55decb96657bdfc935007d Mon Sep 17 00:00:00 2001 From: Khoroshevskyi Date: Wed, 17 Dec 2025 09:52:08 -0500 Subject: [PATCH 02/11] linting, isort, and bedfile file refactoring --- .github/workflows/black.yml | 6 +- .github/workflows/cli-coverage.yml | 4 +- .github/workflows/run-pytest.yml | 2 +- bbconf/bbagent.py | 26 ++-- bbconf/config_parser/bedbaseconfig.py | 62 +++------ bbconf/config_parser/models.py | 6 +- bbconf/db_utils.py | 4 +- bbconf/models/base_models.py | 2 +- bbconf/modules/bedfiles.py | 173 +++++++++++++------------- manual_testing.py | 6 +- tests/test_common.py | 7 +- 11 files changed, 137 insertions(+), 161 deletions(-) diff --git a/.github/workflows/black.yml b/.github/workflows/black.yml index 63e18519..93ef3f6f 100644 --- a/.github/workflows/black.yml +++ b/.github/workflows/black.yml @@ -6,6 +6,6 @@ jobs: lint: runs-on: ubuntu-latest steps: - - uses: actions/checkout@v2 - - uses: actions/setup-python@v2 - - uses: psf/black@20.8b1 + - uses: actions/checkout@v6 + - uses: actions/setup-python@v6 + - uses: psf/black@@stable diff --git a/.github/workflows/cli-coverage.yml b/.github/workflows/cli-coverage.yml index 06a23e4f..5d355745 100644 --- a/.github/workflows/cli-coverage.yml +++ b/.github/workflows/cli-coverage.yml @@ -8,7 +8,7 @@ jobs: cli-coverage-report: strategy: matrix: - python-version: ["3.10"] + python-version: ["3.12"] os: [ ubuntu-latest ] # can't use macOS when using service containers or container jobs runs-on: ${{ matrix.os }} services: @@ -28,7 +28,7 @@ jobs: - uses: actions/setup-python@v5 with: - python-version: '3.10' + python-version: '3.12' - name: Install uv run: pip install uv diff --git a/.github/workflows/run-pytest.yml b/.github/workflows/run-pytest.yml index 876596f3..5b594a08 100644 --- a/.github/workflows/run-pytest.yml +++ b/.github/workflows/run-pytest.yml @@ -12,7 +12,7 @@ jobs: pytest: strategy: matrix: - python-version: ["3.9", "3.12"] + python-version: ["3.10", "3.13"] os: [ubuntu-latest] # can't use macOS when using service containers or container jobs runs-on: ${{ matrix.os }} services: diff --git a/bbconf/bbagent.py b/bbconf/bbagent.py index 59a58b6d..ca8de8c8 100644 --- a/bbconf/bbagent.py +++ b/bbconf/bbagent.py @@ -1,36 +1,36 @@ import logging +import statistics from functools import cached_property from pathlib import Path -from typing import List, Union, Dict -import numpy as np -import statistics +from typing import Dict, List, Union +import numpy as np from sqlalchemy.orm import Session -from sqlalchemy.sql import distinct, func, select, and_, or_ +from sqlalchemy.sql import and_, distinct, func, or_, select from bbconf.config_parser.bedbaseconfig import BedBaseConfig from bbconf.db_utils import ( Bed, BedMetadata, BedSets, + BedStats, + Files, + GeoGsmStatus, License, - UsageBedSetMeta, UsageBedMeta, + UsageBedSetMeta, UsageFiles, UsageSearch, - GeoGsmStatus, - BedStats, - Files, ) from bbconf.models.base_models import ( - StatsReturn, - UsageModel, - FileStats, - UsageStats, AllFilesInfo, - FileInfo, BinValues, + FileInfo, + FileStats, GEOStatistics, + StatsReturn, + UsageModel, + UsageStats, ) from bbconf.modules.bedfiles import BedAgentBedFile from bbconf.modules.bedsets import BedAgentBedSet diff --git a/bbconf/config_parser/bedbaseconfig.py b/bbconf/config_parser/bedbaseconfig.py index 87a76a41..c96edca5 100644 --- a/bbconf/config_parser/bedbaseconfig.py +++ b/bbconf/config_parser/bedbaseconfig.py @@ -1,3 +1,4 @@ +import io import logging import os import warnings @@ -5,28 +6,24 @@ from typing import List, Literal, Union import boto3 +import joblib import qdrant_client -from qdrant_client import QdrantClient, models +import requests import s3fs import yacman import zarr -import requests - -from fastembed import TextEmbedding -from sentence_transformers import SparseEncoder - from botocore.exceptions import BotoCoreError, EndpointConnectionError +from fastembed import TextEmbedding from geniml.region2vec.main import Region2VecExModel from geniml.search import BED2BEDSearchInterface from geniml.search.backends import BiVectorBackend, QdrantBackend from geniml.search.interfaces import BiVectorSearchInterface from geniml.search.query2vec import BED2Vec from pephubclient import PEPHubClient -from zarr import Group as Z_GROUP +from qdrant_client import QdrantClient, models +from sentence_transformers import SparseEncoder from umap import UMAP - -import joblib -import io +from zarr import Group as Z_GROUP from bbconf.config_parser.const import ( S3_BEDSET_PATH_FOLDER, @@ -80,7 +77,7 @@ def __init__(self, config: Union[Path, str], init_ml: bool = True): self.dense_encoder: TextEmbedding = self._init_dense_encoder() self.sparce_encoder: Union[SparseEncoder, None] = self._init_sparce_model() self._umap_encoder: Union[UMAP, None] = self._init_umap_model() - self.r2v_encoder = self._init_r2v_encoder() + self.r2v_encoder: Union[Region2VecExModel, None] = self._init_r2v_encoder() self._init_qdrant_hybrid( qdrant_cl=self.qdrant_client, @@ -97,11 +94,12 @@ def __init__(self, config: Union[Path, str], init_ml: bool = True): ) ) # used for bivec search - self._b2b_search_interface = self._init_b2b_search_interface( + self.b2b_search_interface = self._init_b2b_search_interface( qdrant_file_backend=self.qdrant_file_backend, + region_encoder=self.r2v_encoder, ) - self._bivec_search_interface = self._init_bivec_interface( + self.bivec_search_interface = self._init_bivec_interface( qdrant_file_backend=self.qdrant_file_backend, qdrant_text_backend=self._qdrant_text_backend, text_encoder=self.dense_encoder, @@ -111,8 +109,8 @@ def __init__(self, config: Union[Path, str], init_ml: bool = True): "Skipping initialization of ML models, init_ml parameter set to False." ) self.r2v_encoder = None - self._b2b_search_interface = None - self._bivec_search_interface = None + self.b2b_search_interface = None + self.bivec_search_interface = None self._umap_encoder: Union[UMAP, None] = None self.sparce_encoder = None @@ -161,34 +159,6 @@ def db_engine(self) -> BaseEngine: """ return self._db_engine - @property - def b2bsi(self) -> Union[BED2BEDSearchInterface, None]: - """ - Get bed2bednn object - - :return: bed2bednn object - """ - return self._b2b_search_interface - - @property - def bivec(self) -> BiVectorSearchInterface: - """ - Get bivec search interface object - - :return: bivec search interface object - """ - - return self._bivec_search_interface - - # @property - # def qdrant_engine(self) -> QdrantBackend: - # """ - # Get qdrant engine - # - # :return: qdrant engine - # """ - # return self._qdrant_file_backend - @property def phc(self) -> PEPHubClient: """ @@ -437,7 +407,9 @@ def _init_bivec_interface( return search_interface def _init_b2b_search_interface( - self, qdrant_file_backend: QdrantBackend + self, + qdrant_file_backend: QdrantBackend, + region_encoder: Union[Region2VecExModel, str], ) -> Union[BED2BEDSearchInterface, None]: """ Create Bed 2 BED search interface and return this object @@ -448,7 +420,7 @@ def _init_b2b_search_interface( _LOGGER.info("Initializing search bed 2 bed search interfaces...") return BED2BEDSearchInterface( backend=qdrant_file_backend, - query2vec=BED2Vec(model=self._config.path.region2vec), + query2vec=BED2Vec(model=region_encoder), ) except Exception as e: _LOGGER.error("Error in creating BED2BEDSearchInterface object: " + str(e)) diff --git a/bbconf/config_parser/models.py b/bbconf/config_parser/models.py index fcfaedbd..b0786dd3 100644 --- a/bbconf/config_parser/models.py +++ b/bbconf/config_parser/models.py @@ -13,16 +13,16 @@ DEFAULT_PEPHUB_NAME, DEFAULT_PEPHUB_NAMESPACE, DEFAULT_PEPHUB_TAG, + DEFAULT_QDRANT_BIVEC_COLLECTION_NAME, DEFAULT_QDRANT_FILE_COLLECTION_NAME, + DEFAULT_QDRANT_HYBRID_COLLECTION_NAME, DEFAULT_QDRANT_PORT, - DEFAULT_QDRANT_BIVEC_COLLECTION_NAME, DEFAULT_REGION2_VEC_MODEL, DEFAULT_S3_BUCKET, DEFAULT_SERVER_HOST, DEFAULT_SERVER_PORT, - DEFAULT_TEXT2VEC_MODEL, DEFAULT_SPARSE_MODEL, - DEFAULT_QDRANT_HYBRID_COLLECTION_NAME, + DEFAULT_TEXT2VEC_MODEL, ) _LOGGER = logging.getLogger(__name__) diff --git a/bbconf/db_utils.py b/bbconf/db_utils.py index d49902a8..b2520e94 100644 --- a/bbconf/db_utils.py +++ b/bbconf/db_utils.py @@ -9,12 +9,12 @@ ForeignKey, Result, Select, + String, UniqueConstraint, event, select, - String, ) -from sqlalchemy.dialects.postgresql import JSON, ARRAY +from sqlalchemy.dialects.postgresql import ARRAY, JSON from sqlalchemy.engine import URL, Engine, create_engine from sqlalchemy.event import listens_for from sqlalchemy.exc import IntegrityError, ProgrammingError diff --git a/bbconf/models/base_models.py b/bbconf/models/base_models.py index a7711857..20c1c26f 100644 --- a/bbconf/models/base_models.py +++ b/bbconf/models/base_models.py @@ -1,5 +1,5 @@ -from typing import List, Optional, Union, Dict import datetime +from typing import Dict, List, Optional, Union from pydantic import BaseModel, ConfigDict, Field diff --git a/bbconf/modules/bedfiles.py b/bbconf/modules/bedfiles.py index cf83775c..4cfbf588 100644 --- a/bbconf/modules/bedfiles.py +++ b/bbconf/modules/bedfiles.py @@ -5,22 +5,19 @@ import numpy as np from geniml.bbclient import BBClient - -# from geniml.io import RegionSet from geniml.search.backends import QdrantBackend from gtars.models import RegionSet as GRegionSet from pephubclient.exceptions import ResponseError from pydantic import BaseModel +from qdrant_client import models from qdrant_client.http.models import PointStruct -from qdrant_client.models import Distance, PointIdsList, VectorParams -from sqlalchemy import and_, delete, func, or_, select, cast +from qdrant_client.models import PointIdsList, QueryResponse +from sqlalchemy import and_, cast, delete, func, or_, select +from sqlalchemy.dialects import postgresql from sqlalchemy.exc import IntegrityError from sqlalchemy.orm import Session, aliased from sqlalchemy.orm.attributes import flag_modified -from sqlalchemy.dialects import postgresql from tqdm import tqdm -from qdrant_client import models -from qdrant_client.models import QueryResponse from bbconf.config_parser.bedbaseconfig import BedBaseConfig from bbconf.const import DEFAULT_LICENSE, PKG_NAME, ZARR_TOKENIZED_FOLDER @@ -86,7 +83,7 @@ def __init__(self, config: BedBaseConfig, bbagent_obj=None): self._sa_engine = config.db_engine.engine self._db_engine = config.db_engine self._boto3_client = config.boto3_client - self._config = config + self.config = config self.bb_agent = bbagent_obj def get(self, identifier: str, full: bool = False) -> BedMetadataAll: @@ -117,7 +114,7 @@ def get(self, identifier: str, full: bool = False) -> BedMetadataAll: FileModel( **result.__dict__, object_id=f"bed.{identifier}.{result.name}", - access_methods=self._config.construct_access_method_list( + access_methods=self.config.construct_access_method_list( result.path ), ), @@ -131,7 +128,7 @@ def get(self, identifier: str, full: bool = False) -> BedMetadataAll: FileModel( **result.__dict__, object_id=f"bed.{identifier}.{result.name}", - access_methods=self._config.construct_access_method_list( + access_methods=self.config.construct_access_method_list( result.path ), ), @@ -167,10 +164,10 @@ def get(self, identifier: str, full: bool = False) -> BedMetadataAll: try: if full: bed_metadata = BedPEPHubRestrict( - **self._config.phc.sample.get( - namespace=self._config.config.phc.namespace, - name=self._config.config.phc.name, - tag=self._config.config.phc.tag, + **self.config.phc.sample.get( + namespace=self.config.config.phc.namespace, + name=self.config.config.phc.name, + tag=self.config.config.phc.tag, sample_name=identifier, ) ) @@ -246,7 +243,7 @@ def get_plots(self, identifier: str) -> BedPlots: FileModel( **result.__dict__, object_id=f"bed.{identifier}.{result.name}", - access_methods=self._config.construct_access_method_list( + access_methods=self.config.construct_access_method_list( result.path ), ), @@ -268,8 +265,8 @@ def get_neighbours( if not self.exists(identifier): raise BEDFileNotFoundError(f"Bed file with id: {identifier} not found.") s = identifier - results = self._config.qdrant_file_backend.qd_client.query_points( - collection_name=self._config.config.qdrant.file_collection, + results = self.config.qdrant_file_backend.qd_client.query_points( + collection_name=self.config.config.qdrant.file_collection, query="-".join([s[:8], s[8:12], s[12:16], s[16:20], s[20:]]), limit=limit, offset=offset, @@ -315,7 +312,7 @@ def get_files(self, identifier: str) -> BedFiles: FileModel( **result.__dict__, object_id=f"bed.{identifier}.{result.name}", - access_methods=self._config.construct_access_method_list( + access_methods=self.config.construct_access_method_list( result.path ), ), @@ -330,10 +327,10 @@ def get_raw_metadata(self, identifier: str) -> BedPEPHub: :return: project metadata """ try: - bed_metadata = self._config.phc.sample.get( - namespace=self._config.config.phc.namespace, - name=self._config.config.phc.name, - tag=self._config.config.phc.tag, + bed_metadata = self.config.phc.sample.get( + namespace=self.config.config.phc.namespace, + name=self.config.config.phc.name, + tag=self.config.config.phc.tag, sample_name=identifier, ) except Exception as e: @@ -386,8 +383,8 @@ def get_embedding(self, identifier: str) -> BedEmbeddingResult: """ if not self.exists(identifier): raise BEDFileNotFoundError(f"Bed file with id: {identifier} not found.") - result = self._config.qdrant_file_backend.qd_client.retrieve( - collection_name=self._config.config.qdrant.file_collection, + result = self.config.qdrant_file_backend.qd_client.retrieve( + collection_name=self.config.config.qdrant.file_collection, ids=[identifier], with_vectors=True, with_payload=True, @@ -625,12 +622,12 @@ def add( # Upload files to s3 if upload_s3: if files: - files = self._config.upload_files_s3( + files = self.config.upload_files_s3( identifier, files=files, base_path=local_path, type="files" ) if plots: - plots = self._config.upload_files_s3( + plots = self.config.upload_files_s3( identifier, files=plots, base_path=local_path, type="plots" ) with Session(self._sa_engine) as session: @@ -931,7 +928,7 @@ def _update_plots( _LOGGER.info("Updating bed file plots..") if plots: - plots = self._config.upload_files_s3( + plots = self.config.upload_files_s3( bed_object.id, files=plots, base_path=local_path, type="plots" ) plots_dict = plots.model_dump( @@ -979,7 +976,7 @@ def _update_files( _LOGGER.info("Updating bed files..") if files: - files = self._config.upload_files_s3( + files = self.config.upload_files_s3( bed_object.id, files=files, base_path=local_path, type="files" ) @@ -1101,16 +1098,16 @@ def delete(self, identifier: str) -> None: self.delete_pephub_sample(identifier) if delete_qdrant: self.delete_qdrant_point(identifier) - self._config.delete_files_s3(files) + self.config.delete_files_s3(files) def upload_pephub(self, identifier: str, metadata: dict, overwrite: bool = False): if not metadata: _LOGGER.warning("No metadata provided. Skipping pephub upload..") return False - self._config.phc.sample.create( - namespace=self._config.config.phc.namespace, - name=self._config.config.phc.name, - tag=self._config.config.phc.tag, + self.config.phc.sample.create( + namespace=self.config.config.phc.namespace, + name=self.config.config.phc.name, + tag=self.config.config.phc.tag, sample_name=identifier, sample_dict=metadata, overwrite=overwrite, @@ -1123,10 +1120,10 @@ def update_pephub( if not metadata: _LOGGER.warning("No metadata provided. Skipping pephub upload..") return None - self._config.phc.sample.update( - namespace=self._config.config.phc.namespace, - name=self._config.config.phc.name, - tag=self._config.config.phc.tag, + self.config.phc.sample.update( + namespace=self.config.config.phc.namespace, + name=self.config.config.phc.name, + tag=self.config.config.phc.tag, sample_name=identifier, sample_dict=metadata, ) @@ -1140,10 +1137,10 @@ def delete_pephub_sample(self, identifier: str): :param identifier: bed file identifier """ try: - self._config.phc.sample.remove( - namespace=self._config.config.phc.namespace, - name=self._config.config.phc.name, - tag=self._config.config.phc.tag, + self.config.phc.sample.remove( + namespace=self.config.config.phc.namespace, + name=self.config.config.phc.name, + tag=self.config.config.phc.tag, sample_name=identifier, ) except ResponseError as e: @@ -1168,12 +1165,12 @@ def upload_file_qdrant( _LOGGER.debug(f"Adding bed file to qdrant. bed_id: {bed_id}") - if not isinstance(self._config.qdrant_file_backend, QdrantBackend): + if not isinstance(self.config.qdrant_file_backend, QdrantBackend): raise QdrantInstanceNotInitializedError("Could not upload file.") bed_embedding = self._embed_file(bed_file) - self._config.qdrant_file_backend.load( + self.config.qdrant_file_backend.load( ids=[bed_id], vectors=bed_embedding, payloads=[{**payload}], @@ -1189,9 +1186,9 @@ def _embed_file(self, bed_file: Union[str, GRegionSet]) -> np.ndarray: :return np array of embeddings """ - if self._config.qdrant_file_backend is None: + if self.config.qdrant_file_backend is None: raise QdrantInstanceNotInitializedError - if not self._config.r2v_encoder: + if not self.config.r2v_encoder: raise BedBaseConfError( "Could not add add region to qdrant. Invalid type, or path. " ) @@ -1208,7 +1205,7 @@ def _embed_file(self, bed_file: Union[str, GRegionSet]) -> np.ndarray: raise BedBaseConfError( "Could not add add region to qdrant. Invalid type, or path. " ) - bed_embedding = np.mean(self._config.r2v_encoder.encode(bed_region_set), axis=0) + bed_embedding = np.mean(self.config.r2v_encoder.encode(bed_region_set), axis=0) vec_dim = bed_embedding.shape[0] return bed_embedding.reshape(1, vec_dim) @@ -1219,11 +1216,11 @@ def _get_umap_file(self, bed_file: Union[str, GRegionSet]) -> np.ndarray: :param bed_file: bed file path or region set """ - if self._config._umap_encoder is None: + if self.config._umap_encoder is None: raise BedBaseConfError("UMAP model is not initialized.") bed_embedding = self._embed_file(bed_file) - bed_umap = self._config._umap_encoder.transform(bed_embedding) + bed_umap = self.config._umap_encoder.transform(bed_embedding) return bed_umap def text_to_bed_search( @@ -1246,7 +1243,9 @@ def text_to_bed_search( """ _LOGGER.info(f"Looking for: {query}") - results = self._config.bivec.query_search(query, limit=limit, offset=offset) + results = self.config.bivec_search_interface.query_search( + query, limit=limit, offset=offset + ) results_list = [] for result in results: result_id = result["id"].replace("-", "") @@ -1267,8 +1266,8 @@ def text_to_bed_search( ) if with_metadata: - count = self._config.qdrant_client.get_collection( - collection_name=self._config.config.qdrant.file_collection + count = self.config.qdrant_client.get_collection( + collection_name=self.config.config.qdrant.file_collection ).points_count else: count = 0 @@ -1294,7 +1293,7 @@ def bed_to_bed_search( :return: BedListSetResults """ - results = self._config.b2bsi.query_search( + results = self.config.b2b_search_interface.query_search( region_set, limit=limit, offset=offset ) results_list = [] @@ -1497,8 +1496,8 @@ def reindex_qdrant(self, batch: int = 100, purge: bool = False) -> None: pbar.set_description( "Uploading points to qdrant using batch..." ) - operation_info = self._config.qdrant_file_backend.qd_client.upsert( - collection_name=self._config.config.qdrant.file_collection, + operation_info = self.config.qdrant_file_backend.qd_client.upsert( + collection_name=self.config.config.qdrant.file_collection, points=points_list, ) pbar.write("Uploaded batch to qdrant.") @@ -1511,8 +1510,8 @@ def reindex_qdrant(self, batch: int = 100, purge: bool = False) -> None: pbar.update(1) _LOGGER.info("Uploading points to qdrant using batches...") - operation_info = self._config.qdrant_file_backend.qd_client.upsert( - collection_name=self._config.config.qdrant.file_collection, + operation_info = self.config.qdrant_file_backend.qd_client.upsert( + collection_name=self.config.config.qdrant.file_collection, points=points_list, ) assert operation_info.status == "completed" @@ -1526,8 +1525,8 @@ def delete_qdrant_point(self, identifier: str) -> None: :return: None """ - result = self._config.qdrant_file_backend.qd_client.delete( - collection_name=self._config.config.qdrant.file_collection, + result = self.config.qdrant_file_backend.qd_client.delete( + collection_name=self.config.config.qdrant.file_collection, points_selector=PointIdsList( points=[identifier], ), @@ -1642,7 +1641,7 @@ def add_tokenized( tokenized_vector=token_vector, overwrite=overwrite, ) - path = os.path.join(f"s3://{self._config.config.s3.bucket}", path) + path = os.path.join(f"s3://{self.config.config.s3.bucket}", path) new_token = TokenizedBed(bed_id=bed_id, universe_id=universe_id, path=path) session.add(new_token) @@ -1665,7 +1664,7 @@ def _add_zarr_s3( :return: zarr path """ - univers_group = self._config.zarr_root.require_group(universe_id) + univers_group = self.config.zarr_root.require_group(universe_id) if not univers_group.get(bed_id): _LOGGER.info("Saving tokenized vector to s3") @@ -1695,7 +1694,7 @@ def get_tokenized(self, bed_id: str, universe_id: str) -> TokenizedBedResponse: if not self.exist_tokenized(bed_id, universe_id): raise TokenizeFileNotExistError("Tokenized file not found in the database.") - univers_group = self._config.zarr_root.require_group(universe_id) + univers_group = self.config.zarr_root.require_group(universe_id) return TokenizedBedResponse( universe_id=universe_id, @@ -1714,7 +1713,7 @@ def delete_tokenized(self, bed_id: str, universe_id: str) -> None: """ if not self.exist_tokenized(bed_id, universe_id): raise TokenizeFileNotExistError("Tokenized file not found in the database.") - univers_group = self._config.zarr_root.require_group(universe_id) + univers_group = self.config.zarr_root.require_group(universe_id) del univers_group[bed_id] @@ -1788,7 +1787,7 @@ def get_tokenized_link( file_path = self._get_tokenized_path(bed_id, universe_id) return TokenizedPathResponse( - endpoint_url=self._config.config.s3.endpoint_url, + endpoint_url=self.config.config.s3.endpoint_url, file_path=file_path, bed_id=bed_id, universe_id=universe_id, @@ -1816,7 +1815,9 @@ def get_missing_plots( t2_alias = aliased(Files) # Define the subquery - subquery = select(t2_alias).where(t2_alias.name == plot_name).subquery() + subquery = ( + select(t2_alias).where(and_(t2_alias.name == plot_name)).subquery() + ) query = ( select(Bed.id) @@ -1871,7 +1872,9 @@ def get_missing_files(self, limit: int = 1000, offset: int = 0) -> List[str]: t2_alias = aliased(Files) # Define the subquery - subquery = select(t2_alias).where(t2_alias.name == "bigbed_file").subquery() + subquery = ( + select(t2_alias).where(and_(t2_alias.name == "bigbed_file")).subquery() + ) query = ( select(Bed.id) @@ -2037,7 +2040,7 @@ def _update_sources( # f"File name {result.name} with summary {result.description}" # ) # - # embeddings_list = list(self._config.dense_encoder.embed(text)) + # embeddings_list = list(self.config.dense_encoder.embed(text)) # # result_list.append( # data = VectorMetadata( # id=result.id, @@ -2068,8 +2071,8 @@ def _update_sources( # pbar.set_description( # "Uploading points to qdrant using batch..." # ) - # operation_info = self._config._qdrant_advanced_engine.upsert( - # collection_name=self._config.config.qdrant.search_collection, + # operation_info = self.config._qdrant_advanced_engine.upsert( + # collection_name=self.config.config.qdrant.search_collection, # points=points, # ) # session.commit() @@ -2080,8 +2083,8 @@ def _update_sources( # pbar.write(f"File: {result.id} successfully indexed.") # pbar.update(1) # - # operation_info = self._config._qdrant_advanced_engine.upsert( - # collection_name=self._config.config.qdrant.search_collection, + # operation_info = self.config._qdrant_advanced_engine.upsert( + # collection_name=self.config.config.qdrant.search_collection, # points=points, # ) # assert operation_info.status == "completed" @@ -2133,10 +2136,10 @@ def reindex_hybrid_search(self, batch: int = 1000, purge: bool = False) -> None: f"File name {result.name} with summary {result.description}" ) - embeddings_list = list(self._config.dense_encoder.embed(text)) + embeddings_list = list(self.config.dense_encoder.embed(text)) - if self._config.sparce_encoder: - sparse_result = self._config.sparce_encoder.encode( + if self.config.sparce_encoder: + sparse_result = self.config.sparce_encoder.encode( text ).coalesce() @@ -2183,8 +2186,8 @@ def reindex_hybrid_search(self, batch: int = 1000, purge: bool = False) -> None: pbar.set_description( "Uploading points to qdrant using batch..." ) - operation_info = self._config.qdrant_client.upsert( - collection_name=self._config.config.qdrant.search_collection, + operation_info = self.config.qdrant_client.upsert( + collection_name=self.config.config.qdrant.search_collection, points=points, ) session.commit() @@ -2195,8 +2198,8 @@ def reindex_hybrid_search(self, batch: int = 1000, purge: bool = False) -> None: pbar.write(f"File: {result.id} successfully indexed.") pbar.update(1) - operation_info = self._config.qdrant_client.upsert( - collection_name=self._config.config.qdrant.search_collection, + operation_info = self.config.qdrant_client.upsert( + collection_name=self.config.config.qdrant.search_collection, points=points, ) assert operation_info.status == "completed" @@ -2244,10 +2247,10 @@ def semantic_search( ) ) - embeddings_list = list(self._config.dense_encoder.embed(query))[0] + embeddings_list = list(self.config.dense_encoder.embed(query))[0] - results: QueryResponse = self._config.qdrant_client.query_points( - collection_name=self._config.config.qdrant.search_collection, + results: QueryResponse = self.config.qdrant_client.query_points( + collection_name=self.config.config.qdrant.search_collection, query=list(embeddings_list), limit=limit, offset=offset, @@ -2332,9 +2335,9 @@ def hybrid_search( ) ) - dense_query = list(list(self._config.dense_encoder.embed(query))[0]) - if self._config.sparce_encoder: - sparse_result = self._config.sparce_encoder.encode(query).coalesce() + dense_query = list(list(self.config.dense_encoder.embed(query))[0]) + if self.config.sparce_encoder: + sparse_result = self.config.sparce_encoder.encode(query).coalesce() sparse_embeddings = models.SparseVector( indices=sparse_result.indices().tolist()[0], values=sparse_result.values().tolist(), @@ -2352,8 +2355,8 @@ def hybrid_search( models.Prefetch(query=dense_query, using="dense", limit=limit), ] - results = self._config.qdrant_client.query_points( - collection_name=self._config.config.qdrant.search_collection, + results = self.config.qdrant_client.query_points( + collection_name=self.config.config.qdrant.search_collection, limit=limit, offset=offset, prefetch=hybrid_query, diff --git a/manual_testing.py b/manual_testing.py index 0ecae33a..ee1fd705 100644 --- a/manual_testing.py +++ b/manual_testing.py @@ -1,14 +1,14 @@ import os +import time +import matplotlib.pyplot as plt +import numpy as np import s3fs import zarr # from dotenv import load_dotenv from geniml.io import RegionSet from gtars.utils import read_tokens_from_gtok -import matplotlib.pyplot as plt -import numpy as np -import time # from gtars.tokenizers import RegionSet diff --git a/tests/test_common.py b/tests/test_common.py index d7921e9b..9e9b029a 100644 --- a/tests/test_common.py +++ b/tests/test_common.py @@ -1,12 +1,13 @@ +import datetime + import pytest from bbconf.const import DEFAULT_LICENSE -from bbconf.models.base_models import UsageModel from bbconf.exceptions import BedBaseConfError -import datetime +from bbconf.models.base_models import UsageModel from .conftest import SERVICE_UNAVAILABLE -from .utils import ContextManagerDBTesting, BED_TEST_ID, BEDSET_TEST_ID +from .utils import BED_TEST_ID, BEDSET_TEST_ID, ContextManagerDBTesting @pytest.mark.skipif(SERVICE_UNAVAILABLE, reason="Database is not available") From b08fa75b6308f35718bff951006d2bcf370999a5 Mon Sep 17 00:00:00 2001 From: Khoroshevskyi Date: Thu, 18 Dec 2025 12:49:33 -0500 Subject: [PATCH 03/11] updated geniml requirements version --- docs/changelog.md | 8 ++++++++ requirements/requirements-all.txt | 2 +- 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/docs/changelog.md b/docs/changelog.md index b0282669..624403c1 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -2,6 +2,14 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html) and [Keep a Changelog](https://keepachangelog.com/en/1.0.0/) format. + +### [0.14.0] - 2025-12-18 +### Fixed: +- Updated qdrant and encoder model initialization. (One connection for qdrant and one initialized object per model) + +### Added: +- Added hybrid semantic search.(dense + sparse search) + ### [0.13.0] - 2025-11-24 ### Added: - Conversion of bedfile to umap from predefined model diff --git a/requirements/requirements-all.txt b/requirements/requirements-all.txt index cbf01e83..56ce87e7 100644 --- a/requirements/requirements-all.txt +++ b/requirements/requirements-all.txt @@ -1,7 +1,7 @@ yacman >= 0.9.1 sqlalchemy >= 2.0.0 gtars >= 0.4.0 -geniml[ml] >= 0.8.2 +geniml[ml] >= 0.8.3 psycopg >= 3.1.15 coloredlogs pydantic >= 2.9.0 From 1bb0ec460e42e495ed4e9ef19ae51af1ea6f054f Mon Sep 17 00:00:00 2001 From: Khoroshevskyi Date: Fri, 19 Dec 2025 13:13:34 -0500 Subject: [PATCH 04/11] fixed add tokenized method --- bbconf/modules/bedfiles.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/bbconf/modules/bedfiles.py b/bbconf/modules/bedfiles.py index 4cfbf588..f2acd36f 100644 --- a/bbconf/modules/bedfiles.py +++ b/bbconf/modules/bedfiles.py @@ -1627,13 +1627,12 @@ def add_tokenized( if self.exist_tokenized(bed_id, universe_id): if not overwrite: - if not overwrite: - raise TokenizeFileExistsError( - "Tokenized file already exists in the database. " - "Set overwrite to True to overwrite it." - ) - else: - self.delete_tokenized(bed_id, universe_id) + raise TokenizeFileExistsError( + "Tokenized file already exists in the database. " + "Set overwrite to True to overwrite it." + ) + else: + self.delete_tokenized(bed_id, universe_id) path = self._add_zarr_s3( bed_id=bed_id, From 153e3802c74b719f191373af47e4e3fd05d688f7 Mon Sep 17 00:00:00 2001 From: Khoroshevskyi Date: Fri, 19 Dec 2025 14:01:27 -0500 Subject: [PATCH 05/11] updated changelog --- docs/changelog.md | 6 ++++++ setup.py | 3 ++- 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/docs/changelog.md b/docs/changelog.md index 624403c1..74b80c0c 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -5,8 +5,14 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm ### [0.14.0] - 2025-12-18 ### Fixed: +- Insertion of tokenized files + +### Changed: - Updated qdrant and encoder model initialization. (One connection for qdrant and one initialized object per model) +### Added: +- Hybrid semantic search: dense + sparse search + ### Added: - Added hybrid semantic search.(dense + sparse search) diff --git a/setup.py b/setup.py index 79956a34..7936fae9 100644 --- a/setup.py +++ b/setup.py @@ -41,9 +41,10 @@ classifiers=[ "Development Status :: 4 - Beta", "License :: OSI Approved :: BSD License", - "Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", "Topic :: Scientific/Engineering :: Bio-Informatics", ], keywords="", From 14319e8f8299b531c927bf5cc239f7cf54bbc1cc Mon Sep 17 00:00:00 2001 From: Khoroshevskyi Date: Fri, 19 Dec 2025 14:05:47 -0500 Subject: [PATCH 06/11] fixed multiple pr comments --- .github/workflows/black.yml | 2 +- bbconf/config_parser/bedbaseconfig.py | 6 +++--- bbconf/modules/bedfiles.py | 8 ++++---- 3 files changed, 8 insertions(+), 8 deletions(-) diff --git a/.github/workflows/black.yml b/.github/workflows/black.yml index 93ef3f6f..d842ae33 100644 --- a/.github/workflows/black.yml +++ b/.github/workflows/black.yml @@ -8,4 +8,4 @@ jobs: steps: - uses: actions/checkout@v6 - uses: actions/setup-python@v6 - - uses: psf/black@@stable + - uses: psf/black@stable diff --git a/bbconf/config_parser/bedbaseconfig.py b/bbconf/config_parser/bedbaseconfig.py index c96edca5..d237e71c 100644 --- a/bbconf/config_parser/bedbaseconfig.py +++ b/bbconf/config_parser/bedbaseconfig.py @@ -75,7 +75,7 @@ def __init__(self, config: Union[Path, str], init_ml: bool = True): if init_ml: self.dense_encoder: TextEmbedding = self._init_dense_encoder() - self.sparce_encoder: Union[SparseEncoder, None] = self._init_sparce_model() + self.sparse_encoder: Union[SparseEncoder, None] = self._init_sparce_model() self._umap_encoder: Union[UMAP, None] = self._init_umap_model() self.r2v_encoder: Union[Region2VecExModel, None] = self._init_r2v_encoder() @@ -112,7 +112,7 @@ def __init__(self, config: Union[Path, str], init_ml: bool = True): self.b2b_search_interface = None self.bivec_search_interface = None self._umap_encoder: Union[UMAP, None] = None - self.sparce_encoder = None + self.sparse_encoder = None self._phc = self._init_pephubclient() self._boto3_client = self._init_boto3_client() @@ -288,7 +288,7 @@ def _init_qdrant_text_backend( return None if not isinstance(dense_encoder, TextEmbedding): _LOGGER.error( - f"Unable to create Qdrant bivec text collection,, dense encoder is None." + f"Unable to create Qdrant bivec text collection, dense encoder is None." ) return None diff --git a/bbconf/modules/bedfiles.py b/bbconf/modules/bedfiles.py index f2acd36f..42a17dc0 100644 --- a/bbconf/modules/bedfiles.py +++ b/bbconf/modules/bedfiles.py @@ -2137,8 +2137,8 @@ def reindex_hybrid_search(self, batch: int = 1000, purge: bool = False) -> None: embeddings_list = list(self.config.dense_encoder.embed(text)) - if self.config.sparce_encoder: - sparse_result = self.config.sparce_encoder.encode( + if self.config.sparse_encoder: + sparse_result = self.config.sparse_encoder.encode( text ).coalesce() @@ -2335,8 +2335,8 @@ def hybrid_search( ) dense_query = list(list(self.config.dense_encoder.embed(query))[0]) - if self.config.sparce_encoder: - sparse_result = self.config.sparce_encoder.encode(query).coalesce() + if self.config.sparse_encoder: + sparse_result = self.config.sparse_encoder.encode(query).coalesce() sparse_embeddings = models.SparseVector( indices=sparse_result.indices().tolist()[0], values=sparse_result.values().tolist(), From 42ce8c13b6533925f3220459fe9b9128a7b45ee7 Mon Sep 17 00:00:00 2001 From: Khoroshevskyi Date: Fri, 19 Dec 2025 14:11:14 -0500 Subject: [PATCH 07/11] commented out old search --- bbconf/modules/bedfiles.py | 180 ++++++++++++++++++------------------- 1 file changed, 90 insertions(+), 90 deletions(-) diff --git a/bbconf/modules/bedfiles.py b/bbconf/modules/bedfiles.py index ead06b8f..04014552 100644 --- a/bbconf/modules/bedfiles.py +++ b/bbconf/modules/bedfiles.py @@ -1190,7 +1190,7 @@ def _embed_file(self, bed_file: Union[str, GRegionSet]) -> np.ndarray: raise QdrantInstanceNotInitializedError if not self.config.r2v_encoder: raise BedBaseConfError( - "Could not add add region to qdrant. Invalid type, or path. " + "Could not add region to qdrant. Invalid type, or path. " ) if isinstance(bed_file, str): @@ -1203,7 +1203,7 @@ def _embed_file(self, bed_file: Union[str, GRegionSet]) -> np.ndarray: bed_region_set = bed_file else: raise BedBaseConfError( - "Could not add add region to qdrant. Invalid type, or path. " + "Could not add region to qdrant. Invalid type, or path. " ) bed_embedding = np.mean(self.config.r2v_encoder.encode(bed_region_set), axis=0) vec_dim = bed_embedding.shape[0] @@ -2106,7 +2106,7 @@ def reindex_hybrid_search(self, batch: int = 1000, purge: bool = False) -> None: select(Bed) .join(BedMetadata, Bed.id == BedMetadata.id) .where(Bed.indexed == False) - .limit(150) + .limit(batch) ) with Session(self._sa_engine) as session: @@ -2206,93 +2206,93 @@ def reindex_hybrid_search(self, batch: int = 1000, purge: bool = False) -> None: return None - def semantic_search( - self, - query: str = "liver", - genome_alias: str = "", - assay: str = "", - limit: int = 100, - offset: int = 0, - with_metadata: bool = True, - ) -> BedListSearchResult: - """ - Run semantic search for bed files using qdrant. - This is not bivec search, but usual qdrant search with embeddings. - - :param query: text query to search for - :param genome_alias: genome alias to filter results - :param assay: filter by assay type - :param limit: number of results to return - :param offset: offset to start from - :param with_metadata: if True, metadata will be returned in the results. Default is True. - - :return: list of bed file metadata - """ - - should_statement = [] - - if genome_alias: - should_statement.append( - models.FieldCondition( - key="genome_alias", - match=models.MatchValue(value=genome_alias), - ) - ) - if assay: - should_statement.append( - models.FieldCondition( - key="assay", - match=models.MatchValue(value=assay), - ) - ) - - embeddings_list = list(self.config.dense_encoder.embed(query))[0] - - results: QueryResponse = self.config.qdrant_client.query_points( - collection_name=self.config.config.qdrant.search_collection, - query=list(embeddings_list), - limit=limit, - offset=offset, - search_params=models.SearchParams( - exact=True, - ), - # query_filter=models.Filter(should=should_statement) if should_statement else None, - query_filter=( - models.Filter(must=should_statement) if should_statement else None - ), - with_payload=True, - with_vectors=True, - ) - - result_list = [] - for result in results.points: - result_id = result.id.replace("-", "") - - if with_metadata: - metadata = self.get(result_id, full=False) - else: - metadata = None - - result_list.append( - QdrantSearchResult( - id=result_id, - payload=result.payload, - score=result.score, - metadata=metadata, - ) - ) - - if with_metadata: - count = self.bb_agent.get_stats().bedfiles_number - else: - count = 0 - - return BedListSearchResult( - count=count, - limit=limit, - offset=offset, - results=result_list, - ) + # def semantic_search( + # self, + # query: str = "liver", + # genome_alias: str = "", + # assay: str = "", + # limit: int = 100, + # offset: int = 0, + # with_metadata: bool = True, + # ) -> BedListSearchResult: + # """ + # Run semantic search for bed files using qdrant. + # This is not bivec search, but usual qdrant search with embeddings. + # + # :param query: text query to search for + # :param genome_alias: genome alias to filter results + # :param assay: filter by assay type + # :param limit: number of results to return + # :param offset: offset to start from + # :param with_metadata: if True, metadata will be returned in the results. Default is True. + # + # :return: list of bed file metadata + # """ + # + # should_statement = [] + # + # if genome_alias: + # should_statement.append( + # models.FieldCondition( + # key="genome_alias", + # match=models.MatchValue(value=genome_alias), + # ) + # ) + # if assay: + # should_statement.append( + # models.FieldCondition( + # key="assay", + # match=models.MatchValue(value=assay), + # ) + # ) + # + # embeddings_list = list(self.config.dense_encoder.embed(query))[0] + # + # results: QueryResponse = self.config.qdrant_client.query_points( + # collection_name=self.config.config.qdrant.search_collection, + # query=list(embeddings_list), + # limit=limit, + # offset=offset, + # search_params=models.SearchParams( + # exact=True, + # ), + # # query_filter=models.Filter(should=should_statement) if should_statement else None, + # query_filter=( + # models.Filter(must=should_statement) if should_statement else None + # ), + # with_payload=True, + # with_vectors=True, + # ) + # + # result_list = [] + # for result in results.points: + # result_id = result.id.replace("-", "") + # + # if with_metadata: + # metadata = self.get(result_id, full=False) + # else: + # metadata = None + # + # result_list.append( + # QdrantSearchResult( + # id=result_id, + # payload=result.payload, + # score=result.score, + # metadata=metadata, + # ) + # ) + # + # if with_metadata: + # count = self.bb_agent.get_stats().bedfiles_number + # else: + # count = 0 + # + # return BedListSearchResult( + # count=count, + # limit=limit, + # offset=offset, + # results=result_list, + # ) def hybrid_search( self, From 3b5de0edbf6e98b9284a54f84a9e620a1539b2a6 Mon Sep 17 00:00:00 2001 From: Khoroshevskyi Date: Fri, 19 Dec 2025 14:15:56 -0500 Subject: [PATCH 08/11] updated config private attribute --- bbconf/config_parser/bedbaseconfig.py | 4 ++-- bbconf/modules/bedfiles.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/bbconf/config_parser/bedbaseconfig.py b/bbconf/config_parser/bedbaseconfig.py index d237e71c..f60ec76f 100644 --- a/bbconf/config_parser/bedbaseconfig.py +++ b/bbconf/config_parser/bedbaseconfig.py @@ -76,7 +76,7 @@ def __init__(self, config: Union[Path, str], init_ml: bool = True): self.dense_encoder: TextEmbedding = self._init_dense_encoder() self.sparse_encoder: Union[SparseEncoder, None] = self._init_sparce_model() - self._umap_encoder: Union[UMAP, None] = self._init_umap_model() + self.umap_encoder: Union[UMAP, None] = self._init_umap_model() self.r2v_encoder: Union[Region2VecExModel, None] = self._init_r2v_encoder() self._init_qdrant_hybrid( @@ -111,7 +111,7 @@ def __init__(self, config: Union[Path, str], init_ml: bool = True): self.r2v_encoder = None self.b2b_search_interface = None self.bivec_search_interface = None - self._umap_encoder: Union[UMAP, None] = None + self.umap_encoder: Union[UMAP, None] = None self.sparse_encoder = None self._phc = self._init_pephubclient() diff --git a/bbconf/modules/bedfiles.py b/bbconf/modules/bedfiles.py index 04014552..b046a61a 100644 --- a/bbconf/modules/bedfiles.py +++ b/bbconf/modules/bedfiles.py @@ -1216,11 +1216,11 @@ def _get_umap_file(self, bed_file: Union[str, GRegionSet]) -> np.ndarray: :param bed_file: bed file path or region set """ - if self.config._umap_encoder is None: + if self.config.umap_encoder is None: raise BedBaseConfError("UMAP model is not initialized.") bed_embedding = self._embed_file(bed_file) - bed_umap = self.config._umap_encoder.transform(bed_embedding) + bed_umap = self.config.umap_encoder.transform(bed_embedding) return bed_umap def text_to_bed_search( From ac7a8c0dee189fa57f636eb6aafda992a8480c5a Mon Sep 17 00:00:00 2001 From: Khoroshevskyi Date: Fri, 19 Dec 2025 14:46:16 -0500 Subject: [PATCH 09/11] qdrant insertion bug fix --- bbconf/modules/bedfiles.py | 23 +++++++++++++++-------- requirements/requirements-all.txt | 3 ++- 2 files changed, 17 insertions(+), 9 deletions(-) diff --git a/bbconf/modules/bedfiles.py b/bbconf/modules/bedfiles.py index b046a61a..13760679 100644 --- a/bbconf/modules/bedfiles.py +++ b/bbconf/modules/bedfiles.py @@ -2197,11 +2197,12 @@ def reindex_hybrid_search(self, batch: int = 1000, purge: bool = False) -> None: pbar.write(f"File: {result.id} successfully indexed.") pbar.update(1) - operation_info = self.config.qdrant_client.upsert( - collection_name=self.config.config.qdrant.search_collection, - points=points, - ) - assert operation_info.status == "completed" + if points: + operation_info = self.config.qdrant_client.upsert( + collection_name=self.config.config.qdrant.search_collection, + points=points, + ) + assert operation_info.status == "completed" session.commit() return None @@ -2344,14 +2345,20 @@ def hybrid_search( hybrid_query = [ # Dense retrieval: semantic understanding - models.Prefetch(query=dense_query, using="dense", limit=limit), + models.Prefetch( + query=dense_query, using="dense", limit=limit + offset + 100 + ), # Sparse retrieval: exact technical term matching - models.Prefetch(query=sparse_embeddings, using="sparse", limit=limit), + models.Prefetch( + query=sparse_embeddings, using="sparse", limit=limit + offset + 100 + ), ] else: hybrid_query = [ # Dense retrieval: semantic understanding - models.Prefetch(query=dense_query, using="dense", limit=limit), + models.Prefetch( + query=dense_query, using="dense", limit=limit + offset + 100 + ), ] results = self.config.qdrant_client.query_points( diff --git a/requirements/requirements-all.txt b/requirements/requirements-all.txt index 56ce87e7..9c3e7060 100644 --- a/requirements/requirements-all.txt +++ b/requirements/requirements-all.txt @@ -1,6 +1,6 @@ yacman >= 0.9.1 sqlalchemy >= 2.0.0 -gtars >= 0.4.0 +gtars >= 0.5.2 geniml[ml] >= 0.8.3 psycopg >= 3.1.15 coloredlogs @@ -15,3 +15,4 @@ s3fs >= 2024.3.1 pandas >= 2.0.0 pybiocfilecache == 0.6.1 umap-learn >= 0.5.8 +qdrant_client >= 1.16.1 From f5eac15b5e428e3e97f73de1efd529b162fcdde4 Mon Sep 17 00:00:00 2001 From: Khoroshevskyi Date: Fri, 19 Dec 2025 16:56:51 -0500 Subject: [PATCH 10/11] renamed search collection to hybrid in config --- bbconf/config_parser/bedbaseconfig.py | 2 +- bbconf/config_parser/models.py | 2 +- bbconf/modules/bedfiles.py | 12 ++++++------ 3 files changed, 8 insertions(+), 8 deletions(-) diff --git a/bbconf/config_parser/bedbaseconfig.py b/bbconf/config_parser/bedbaseconfig.py index f60ec76f..a37c5bcc 100644 --- a/bbconf/config_parser/bedbaseconfig.py +++ b/bbconf/config_parser/bedbaseconfig.py @@ -328,7 +328,7 @@ def _init_qdrant_hybrid( return None dimensions = int(dense_encoder.get_embedding_size(self._config.path.text2vec)) - collection_name = self.config.qdrant.search_collection + collection_name = self.config.qdrant.hybrid_collection try: if not qdrant_cl.collection_exists(collection_name): diff --git a/bbconf/config_parser/models.py b/bbconf/config_parser/models.py index b0786dd3..3fe3fbe2 100644 --- a/bbconf/config_parser/models.py +++ b/bbconf/config_parser/models.py @@ -56,7 +56,7 @@ class ConfigQdrant(BaseModel): api_key: Optional[str] = None file_collection: str = DEFAULT_QDRANT_FILE_COLLECTION_NAME text_collection: Optional[str] = DEFAULT_QDRANT_BIVEC_COLLECTION_NAME - search_collection: Optional[str] = DEFAULT_QDRANT_HYBRID_COLLECTION_NAME + hybrid_collection: Optional[str] = DEFAULT_QDRANT_HYBRID_COLLECTION_NAME class ConfigServer(BaseModel): diff --git a/bbconf/modules/bedfiles.py b/bbconf/modules/bedfiles.py index 13760679..89897fd2 100644 --- a/bbconf/modules/bedfiles.py +++ b/bbconf/modules/bedfiles.py @@ -2071,7 +2071,7 @@ def _update_sources( # "Uploading points to qdrant using batch..." # ) # operation_info = self.config._qdrant_advanced_engine.upsert( - # collection_name=self.config.config.qdrant.search_collection, + # collection_name=self.config.config.qdrant.hybrid_collection, # points=points, # ) # session.commit() @@ -2083,7 +2083,7 @@ def _update_sources( # pbar.update(1) # # operation_info = self.config._qdrant_advanced_engine.upsert( - # collection_name=self.config.config.qdrant.search_collection, + # collection_name=self.config.config.qdrant.hybrid_collection, # points=points, # ) # assert operation_info.status == "completed" @@ -2186,7 +2186,7 @@ def reindex_hybrid_search(self, batch: int = 1000, purge: bool = False) -> None: "Uploading points to qdrant using batch..." ) operation_info = self.config.qdrant_client.upsert( - collection_name=self.config.config.qdrant.search_collection, + collection_name=self.config.config.qdrant.hybrid_collection, points=points, ) session.commit() @@ -2199,7 +2199,7 @@ def reindex_hybrid_search(self, batch: int = 1000, purge: bool = False) -> None: if points: operation_info = self.config.qdrant_client.upsert( - collection_name=self.config.config.qdrant.search_collection, + collection_name=self.config.config.qdrant.hybrid_collection, points=points, ) assert operation_info.status == "completed" @@ -2250,7 +2250,7 @@ def reindex_hybrid_search(self, batch: int = 1000, purge: bool = False) -> None: # embeddings_list = list(self.config.dense_encoder.embed(query))[0] # # results: QueryResponse = self.config.qdrant_client.query_points( - # collection_name=self.config.config.qdrant.search_collection, + # collection_name=self.config.config.qdrant.hybrid_collection, # query=list(embeddings_list), # limit=limit, # offset=offset, @@ -2362,7 +2362,7 @@ def hybrid_search( ] results = self.config.qdrant_client.query_points( - collection_name=self.config.config.qdrant.search_collection, + collection_name=self.config.config.qdrant.hybrid_collection, limit=limit, offset=offset, prefetch=hybrid_query, From 40f2123ae955d664786230113e1ca75902af9f66 Mon Sep 17 00:00:00 2001 From: Khoroshevskyi Date: Fri, 19 Dec 2025 18:27:10 -0500 Subject: [PATCH 11/11] fixed #97 --- bbconf/modules/bedfiles.py | 38 +++++++++++++++++++++++--------------- 1 file changed, 23 insertions(+), 15 deletions(-) diff --git a/bbconf/modules/bedfiles.py b/bbconf/modules/bedfiles.py index 89897fd2..1db3a246 100644 --- a/bbconf/modules/bedfiles.py +++ b/bbconf/modules/bedfiles.py @@ -11,6 +11,7 @@ from pydantic import BaseModel from qdrant_client import models from qdrant_client.http.models import PointStruct +from qdrant_client.http.exceptions import UnexpectedResponse from qdrant_client.models import PointIdsList, QueryResponse from sqlalchemy import and_, cast, delete, func, or_, select from sqlalchemy.dialects import postgresql @@ -265,23 +266,30 @@ def get_neighbours( if not self.exists(identifier): raise BEDFileNotFoundError(f"Bed file with id: {identifier} not found.") s = identifier - results = self.config.qdrant_file_backend.qd_client.query_points( - collection_name=self.config.config.qdrant.file_collection, - query="-".join([s[:8], s[8:12], s[12:16], s[16:20], s[20:]]), - limit=limit, - offset=offset, - ) - result_list = [] - for result in results.points: - result_id = result.id.replace("-", "") - result_list.append( - QdrantSearchResult( - id=result_id, - payload=result.payload, - score=result.score, - metadata=self.get(result_id, full=False), + try: + results = self.config.qdrant_file_backend.qd_client.query_points( + collection_name=self.config.config.qdrant.file_collection, + query="-".join([s[:8], s[8:12], s[12:16], s[16:20], s[20:]]), + limit=limit, + offset=offset, + ) + result_list = [] + for result in results.points: + result_id = result.id.replace("-", "") + result_list.append( + QdrantSearchResult( + id=result_id, + payload=result.payload, + score=result.score, + metadata=self.get(result_id, full=False), + ) ) + except UnexpectedResponse as err: + _LOGGER.error( + f"Qdrant request failed. Error: {err}. Returning empty result set." ) + result_list = [] + return BedListSearchResult( count=self.bb_agent.get_stats().bedfiles_number, limit=limit,