diff --git a/.github/workflows/black.yml b/.github/workflows/black.yml index 63e18519..d842ae33 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/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/.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/_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/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 d3244b33..a37c5bcc 100644 --- a/bbconf/config_parser/bedbaseconfig.py +++ b/bbconf/config_parser/bedbaseconfig.py @@ -1,34 +1,34 @@ +import io import logging import os import warnings from pathlib import Path from typing import List, Literal, Union -import numpy as np import boto3 +import joblib import qdrant_client -from qdrant_client import QdrantClient, models +import requests import s3fs import yacman import zarr 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, 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 +64,55 @@ 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.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() + + 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, + region_encoder=self.r2v_encoder, + ) + + 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.sparse_encoder = None self._phc = self._init_pephubclient() self._boto3_client = self._init_boto3_client() @@ -128,43 +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._b2bsi - - @property - def r2v(self) -> Region2VecExModel: - """ - Get region2vec object - - :return: region2vec object - """ - return self._r2v - - @property - def bivec(self) -> BiVectorSearchInterface: - """ - Get bivec search interface object - - :return: bivec search interface object - """ - - return self._bivec - - @property - def qdrant_engine(self) -> QdrantBackend: - """ - Get qdrant engine - - :return: qdrant engine - """ - return self._qdrant_engine - @property def phc(self) -> PEPHubClient: """ @@ -226,81 +220,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 text advanced engine...") + _LOGGER.info("Initializing qdrant sparse collection...") - 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 + + dimensions = int(dense_encoder.get_embedding_size(self._config.path.text2vec)) + collection_name = self.config.qdrant.hybrid_collection - if not qdrant_cl.collection_exists(COLLECTION_NAME): + 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,56 +360,67 @@ 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, + region_encoder: Union[Region2VecExModel, str], + ) -> 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, - query2vec=BED2Vec(model=self._config.path.region2vec), + backend=qdrant_file_backend, + query2vec=BED2Vec(model=region_encoder), ) except Exception as e: _LOGGER.error("Error in creating BED2BEDSearchInterface object: " + str(e)) @@ -370,56 +430,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 +484,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 +512,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 +648,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..3fe3fbe2 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_BIVEC_COLLECTION_NAME, + DEFAULT_QDRANT_FILE_COLLECTION_NAME, + DEFAULT_QDRANT_HYBRID_COLLECTION_NAME, DEFAULT_QDRANT_PORT, - DEFAULT_QDRANT_TEXT_COLLECTION_NAME, DEFAULT_REGION2_VEC_MODEL, DEFAULT_S3_BUCKET, DEFAULT_SERVER_HOST, DEFAULT_SERVER_PORT, + DEFAULT_SPARSE_MODEL, DEFAULT_TEXT2VEC_MODEL, - DEFAULT_QDRANT_SEARCH_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 + hybrid_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/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 70633c68..1db3a246 100644 --- a/bbconf/modules/bedfiles.py +++ b/bbconf/modules/bedfiles.py @@ -5,22 +5,20 @@ 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.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 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 fastembed import TextEmbedding -from qdrant_client import models from bbconf.config_parser.bedbaseconfig import BedBaseConfig from bbconf.const import DEFAULT_LICENSE, PKG_NAME, ZARR_TOKENIZED_FOLDER @@ -85,14 +83,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.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. @@ -121,7 +115,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 ), ), @@ -135,7 +129,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 ), ), @@ -171,10 +165,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, ) ) @@ -250,7 +244,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 ), ), @@ -272,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._qdrant_engine.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, @@ -319,7 +320,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 ), ), @@ -334,10 +335,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: @@ -390,8 +391,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._qdrant_engine.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, @@ -629,12 +630,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: @@ -935,7 +936,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( @@ -983,7 +984,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" ) @@ -1105,16 +1106,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, @@ -1127,10 +1128,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, ) @@ -1144,10 +1145,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: @@ -1172,12 +1173,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,11 +1194,11 @@ 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. " + "Could not add region to qdrant. Invalid type, or path. " ) if isinstance(bed_file, str): @@ -1210,9 +1211,9 @@ 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.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 +1224,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( @@ -1250,7 +1251,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("-", "") @@ -1271,8 +1274,8 @@ def text_to_bed_search( ) if with_metadata: - count = self._config._qdrant_advanced_engine.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 @@ -1298,7 +1301,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 = [] @@ -1501,8 +1504,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_engine.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.") @@ -1515,8 +1518,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_engine.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" @@ -1530,22 +1533,13 @@ def delete_qdrant_point(self, identifier: str) -> None: :return: None """ - result = self._config.qdrant_engine.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], ), ) - 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: """ @@ -1641,13 +1635,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, @@ -1655,7 +1648,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) @@ -1678,7 +1671,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") @@ -1708,7 +1701,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, @@ -1727,7 +1720,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] @@ -1801,7 +1794,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, @@ -1829,7 +1822,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) @@ -1884,7 +1879,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) @@ -2006,7 +2003,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.hybrid_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.hybrid_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 +2114,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(batch) ) with Session(self._sa_engine) as session: @@ -2050,7 +2143,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.sparse_encoder: + sparse_result = self.config.sparse_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 +2182,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,8 +2193,8 @@ 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( - collection_name=self._config.config.qdrant.search_collection, + operation_info = self.config.qdrant_client.upsert( + collection_name=self.config.config.qdrant.hybrid_collection, points=points, ) session.commit() @@ -2093,16 +2205,105 @@ 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( - 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.hybrid_collection, + points=points, + ) + assert operation_info.status == "completed" session.commit() return None - def semantic_search( + # 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.hybrid_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, query: str = "liver", genome_alias: str = "", @@ -2113,7 +2314,7 @@ def semantic_search( ) -> BedListSearchResult: """ Run semantic search for bed files using qdrant. - This is not bivec search, but usual qdrant search with embeddings. + 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 @@ -2125,43 +2326,67 @@ def semantic_search( :return: list of bed file metadata """ - should_statement = [] + must_statement = [] if genome_alias: - should_statement.append( + must_statement.append( models.FieldCondition( key="genome_alias", match=models.MatchValue(value=genome_alias), ) ) if assay: - should_statement.append( + must_statement.append( models.FieldCondition( key="assay", match=models.MatchValue(value=assay), ) ) - embeddings_list = list(self._embedding_model.embed(query))[0] + dense_query = list(list(self.config.dense_encoder.embed(query))[0]) + 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(), + ) + + hybrid_query = [ + # Dense retrieval: semantic understanding + 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 + offset + 100 + ), + ] + else: + hybrid_query = [ + # Dense retrieval: semantic understanding + models.Prefetch( + query=dense_query, using="dense", limit=limit + offset + 100 + ), + ] - results = self._config._qdrant_advanced_engine.search( - collection_name=self._config.config.qdrant.search_collection, - query_vector=list(embeddings_list), + results = self.config.qdrant_client.query_points( + collection_name=self.config.config.qdrant.hybrid_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(should=should_statement) if should_statement else None, query_filter=( - models.Filter(must=should_statement) if should_statement else None + models.Filter(must=must_statement) if must_statement else None ), - with_payload=True, - with_vectors=True, ) result_list = [] - for result in results: + for result in results.points: result_id = result.id.replace("-", "") if with_metadata: @@ -2205,7 +2430,7 @@ def search_external_file(self, source: str, accession: str) -> BedListSearchResu f"Source {source} is not supported. Supported sources are: 'geo', 'encode'." ) - if source == "geo" and accession.startswith("gse"): + if source == "geo" and accession.upper().startswith("GSE"): statement = ( select(Bed) .join(BedMetadata, Bed.id == BedMetadata.id) diff --git a/docs/changelog.md b/docs/changelog.md index b0282669..74b80c0c 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -2,6 +2,20 @@ 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: +- 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) + ### [0.13.0] - 2025-11-24 ### Added: - Conversion of bedfile to umap from predefined model diff --git a/manual_testing.py b/manual_testing.py index 2fc2205d..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 @@ -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..9c3e7060 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 +gtars >= 0.5.2 +geniml[ml] >= 0.8.3 psycopg >= 3.1.15 coloredlogs pydantic >= 2.9.0 @@ -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 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="", 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 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")