Skip to content

Config

helical.models.caduceus.CaduceusConfig

Source code in helical/models/caduceus/caduceus_config.py 
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
class CaduceusConfig:
    def __init__(
        self,
        model_name: Literal[
            "caduceus-ph-4L-seqlen-1k-d118",
            "caduceus-ph-4L-seqlen-1k-d256",
            "caduceus-ph-16L-seqlen-131k-d256",
            "caduceus-ps-4L-seqlen-1k-d118",
            "caduceus-ps-4L-seqlen-1k-d256",
            "caduceus-ps-16L-seqlen-131k-d256",
        ] = "caduceus-ph-4L-seqlen-1k-d118",
        batch_size: int = 5,
        pooling_strategy: Literal["mean", "max", "last", "first"] = "mean",
        nproc: int = 1,
    ):
        """
        Configuration class to use the Caduceus Model.

        Parameters
        ----------
        model_name : Literal["caduceus-ph-4L-seqlen-1k-d118", "caduceus-ph-4L-seqlen-1k-d256", "caduceus-ph-16L-seqlen-131k-d256", "caduceus-ps-4L-seqlen-1k-d118", "caduceus-ps-4L-seqlen-1k-d256", "caduceus-ps-16L-seqlen-131k-d256"], optional, default="caduceus-ph-4L-seqlen-1k-d118"
            The model configuration name to use.
        batch_size : int, optional, default=5
            The batch size. This will be used in all the model operations.
        pooling_strategy : Literal["mean", "max", "last", "first"], optional, default="mean"
            The pooling strategy to use. This will be used in all the model operations.
        nproc : int, optional, default=1
            The number of processes to use for data processing.

        Notes
        ----------
        Caduceus uses the mamba_ssm package which requires CUDA.
        """

        if pooling_strategy not in ["mean", "max", "last", "first"]:
            error = f"Pooling strategy {pooling_strategy} not found in available strategies: ['mean', 'max', 'last', 'first']"
            LOGGER.error(error)
            raise ValueError(error)

        model_map = {
            "caduceus-ph-4L-seqlen-1k-d118": {
                "input_size": 1024,
                "embedding_size": 118,
            },
            "caduceus-ph-4L-seqlen-1k-d256": {
                "input_size": 1024,
                "embedding_size": 256,
            },
            "caduceus-ph-16L-seqlen-131k-d256": {
                "input_size": 1024 * 128,
                "embedding_size": 256,
            },
            "caduceus-ps-4L-seqlen-1k-d118": {
                "input_size": 1024,
                "embedding_size": 118,
            },
            "caduceus-ps-4L-seqlen-1k-d256": {
                "input_size": 1024,
                "embedding_size": 256,
            },
            "caduceus-ps-16L-seqlen-131k-d256": {
                "input_size": 1024 * 128,
                "embedding_size": 256,
            },
        }

        if model_name not in model_map.keys():
            error = f"Model name {model_name} not found in available models: {model_map.keys()}"
            LOGGER.error(error)
            raise ValueError(error)

        self.list_of_files_to_download = [
            f"caduceus/{model_name}/model.safetensors",
            f"caduceus/{model_name}/config.json",
        ]

        self.files_config = {
            "model_files_dir": Path(CACHE_DIR_HELICAL, "caduceus", model_name),
        }

        self.config = {
            "input_size": model_map[model_name]["input_size"],
            "embedding_size": model_map[model_name]["embedding_size"],
            "model_name": model_name,
            "pooling_strategy": pooling_strategy,
            "batch_size": batch_size,
            "device": "cuda",
            "nproc": nproc,
        }

__init__(model_name='caduceus-ph-4L-seqlen-1k-d118', batch_size=5, pooling_strategy='mean', nproc=1)

Configuration class to use the Caduceus Model.

Parameters:

Name Type Description Default
model_name Literal['caduceus-ph-4L-seqlen-1k-d118', 'caduceus-ph-4L-seqlen-1k-d256', 'caduceus-ph-16L-seqlen-131k-d256', 'caduceus-ps-4L-seqlen-1k-d118', 'caduceus-ps-4L-seqlen-1k-d256', 'caduceus-ps-16L-seqlen-131k-d256']

The model configuration name to use.

 "caduceus-ph-4L-seqlen-1k-d118"
batch_size int

The batch size. This will be used in all the model operations.

 5
pooling_strategy Literal['mean', 'max', 'last', 'first']

The pooling strategy to use. This will be used in all the model operations.

 "mean"
nproc int

The number of processes to use for data processing.

 1
Notes

Caduceus uses the mamba_ssm package which requires CUDA.

Source code in helical/models/caduceus/caduceus_config.py 
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
def __init__(
    self,
    model_name: Literal[
        "caduceus-ph-4L-seqlen-1k-d118",
        "caduceus-ph-4L-seqlen-1k-d256",
        "caduceus-ph-16L-seqlen-131k-d256",
        "caduceus-ps-4L-seqlen-1k-d118",
        "caduceus-ps-4L-seqlen-1k-d256",
        "caduceus-ps-16L-seqlen-131k-d256",
    ] = "caduceus-ph-4L-seqlen-1k-d118",
    batch_size: int = 5,
    pooling_strategy: Literal["mean", "max", "last", "first"] = "mean",
    nproc: int = 1,
):
    """
    Configuration class to use the Caduceus Model.

    Parameters
    ----------
    model_name : Literal["caduceus-ph-4L-seqlen-1k-d118", "caduceus-ph-4L-seqlen-1k-d256", "caduceus-ph-16L-seqlen-131k-d256", "caduceus-ps-4L-seqlen-1k-d118", "caduceus-ps-4L-seqlen-1k-d256", "caduceus-ps-16L-seqlen-131k-d256"], optional, default="caduceus-ph-4L-seqlen-1k-d118"
        The model configuration name to use.
    batch_size : int, optional, default=5
        The batch size. This will be used in all the model operations.
    pooling_strategy : Literal["mean", "max", "last", "first"], optional, default="mean"
        The pooling strategy to use. This will be used in all the model operations.
    nproc : int, optional, default=1
        The number of processes to use for data processing.

    Notes
    ----------
    Caduceus uses the mamba_ssm package which requires CUDA.
    """

    if pooling_strategy not in ["mean", "max", "last", "first"]:
        error = f"Pooling strategy {pooling_strategy} not found in available strategies: ['mean', 'max', 'last', 'first']"
        LOGGER.error(error)
        raise ValueError(error)

    model_map = {
        "caduceus-ph-4L-seqlen-1k-d118": {
            "input_size": 1024,
            "embedding_size": 118,
        },
        "caduceus-ph-4L-seqlen-1k-d256": {
            "input_size": 1024,
            "embedding_size": 256,
        },
        "caduceus-ph-16L-seqlen-131k-d256": {
            "input_size": 1024 * 128,
            "embedding_size": 256,
        },
        "caduceus-ps-4L-seqlen-1k-d118": {
            "input_size": 1024,
            "embedding_size": 118,
        },
        "caduceus-ps-4L-seqlen-1k-d256": {
            "input_size": 1024,
            "embedding_size": 256,
        },
        "caduceus-ps-16L-seqlen-131k-d256": {
            "input_size": 1024 * 128,
            "embedding_size": 256,
        },
    }

    if model_name not in model_map.keys():
        error = f"Model name {model_name} not found in available models: {model_map.keys()}"
        LOGGER.error(error)
        raise ValueError(error)

    self.list_of_files_to_download = [
        f"caduceus/{model_name}/model.safetensors",
        f"caduceus/{model_name}/config.json",
    ]

    self.files_config = {
        "model_files_dir": Path(CACHE_DIR_HELICAL, "caduceus", model_name),
    }

    self.config = {
        "input_size": model_map[model_name]["input_size"],
        "embedding_size": model_map[model_name]["embedding_size"],
        "model_name": model_name,
        "pooling_strategy": pooling_strategy,
        "batch_size": batch_size,
        "device": "cuda",
        "nproc": nproc,
    }