ferritin_plms/esm/utils/constants/
models.rs

1// Model names
2pub const ESM3_OPEN_SMALL: &str = "esm3_sm_open_v1";
3pub const ESM3_OPEN_SMALL_ALIAS_1: &str = "esm3-open-2024-03";
4pub const ESM3_OPEN_SMALL_ALIAS_2: &str = "esm3-sm-open-v1";
5pub const ESM3_OPEN_SMALL_ALIAS_3: &str = "esm3-open";
6pub const ESM3_STRUCTURE_ENCODER_V0: &str = "esm3_structure_encoder_v0";
7pub const ESM3_STRUCTURE_DECODER_V0: &str = "esm3_structure_decoder_v0";
8pub const ESM3_FUNCTION_DECODER_V0: &str = "esm3_function_decoder_v0";
9pub const ESMC_600M: &str = "esmc_600m";
10pub const ESMC_300M: &str = "esmc_300m";
11
12pub fn model_is_locally_supported(x: &str) -> bool {
13    matches!(
14        x,
15        ESM3_OPEN_SMALL
16            | ESM3_OPEN_SMALL_ALIAS_1
17            | ESM3_OPEN_SMALL_ALIAS_2
18            | ESM3_OPEN_SMALL_ALIAS_3
19    )
20}
21
22pub fn normalize_model_name(x: &str) -> &str {
23    if matches!(
24        x,
25        ESM3_OPEN_SMALL_ALIAS_1 | ESM3_OPEN_SMALL_ALIAS_2 | ESM3_OPEN_SMALL_ALIAS_3
26    ) {
27        ESM3_OPEN_SMALL
28    } else {
29        x
30    }
31}