check in zero shot cl eval.

Author: Jacoberts

backend/src/folde/campaign.py

@@ -461,11 +461,9 @@ def simulate_campaign(
             cache_key
         ]
         naturalness_ensemble_df = entire_naturalness_df[
-            (
-                "wt_marginal"
-                if model_config.few_shot_naturalness_column is None
-                else model_config.few_shot_naturalness_column
-            )
+            ["wt_marginal"]
+            if model_config.naturalness_columns is None
+            else model_config.naturalness_columns
         ]
         embedding_series = embedding_df[
             "embedding" if model_config.embedding_column is None else model_config.embedding_column

backend/src/folde/types.py

@@ -7,14 +7,17 @@ class FolDEModelConfig(BaseModel):
     name: str
     data_split_mode: Optional[str] = None  # Can be "1-VS-REST" or "2-VS-REST" etc.
     naturalness_model_id: str
+    naturalness_columns: Optional[List[str]] = None
     embedding_model_id: str
     embedding_column: Optional[str] = None
     zero_shot_model_name: str
     zero_shot_model_params: Dict[str, Any]
     few_shot_model_name: str
-    few_shot_naturalness_column: Optional[str] = None
     few_shot_model_params: Dict[str, Any]
 
+    # DEPRECATED.
+    few_shot_naturalness_column: Optional[str] = None
+
 
 class MutantMetrics(BaseModel):
     """Stores dense information about each mutants tested in the simulation."""

backend/src/folde/zero_shot_models.py

@@ -113,20 +113,7 @@ def get_top_n(
         self.selection_debug_info = {"sorts": {}}
 
         if self.decision_mode == "constantliar" or self.decision_mode == "krigingbeliever":
-            if self.lie_noise_stddev_multiplier is not None:
-                lie_noise_stddev_multiplier = self.lie_noise_stddev_multiplier
-            elif self.lie_noise_stddev_multiplier_schedule is not None:
-                lie_noise_stddev_multiplier = self.lie_noise_stddev_multiplier_schedule[0]
-                self.lie_noise_stddev_multiplier_schedule = (
-                    self.lie_noise_stddev_multiplier_schedule[1:]
-                )
-                logging.info(
-                    f"Using lie_noise_stddev_multiplier {lie_noise_stddev_multiplier} for this round and {self.lie_noise_stddev_multiplier_schedule} remaining."
-                )
-            else:
-                raise ValueError(
-                    "Either lie_noise_stddev_multiplier or lie_noise_stddev_multiplier_schedule must be set."
-                )
+            assert self.lie_noise_stddev_multiplier is not None
 
             ensemble_scores = get_consensus_scores(ensemble_of_predictions, "mean")
             pred_df = {
@@ -142,7 +129,7 @@ def get_top_n(
                 cl_considerations.to_numpy(),
                 cl_considerations.index.to_numpy(),
                 n,
-                lie_noise_stddev_multiplier=lie_noise_stddev_multiplier,
+                lie_noise_stddev_multiplier=self.lie_noise_stddev_multiplier,
                 choice_of_baseline="min" if self.decision_mode == "constantliar" else "mean",
                 ucb_beta=self.ucb_beta if self.ucb_beta is not None else 0.0,
             )
@@ -229,41 +216,37 @@ def __init__(self, **kwargs):
         """
         super().__init__(**kwargs)
         self.is_pretrained = False
-        self.single_mutant_naturalness_ensemble: Optional[pd.DataFrame] = None
+        self.single_mutant_naturalness_df: Optional[pd.DataFrame] = None
         self.knns: Dict[str, KNeighborsRegressor] = {}
 
     def pretrain(
         self,
         naturalness_df: pd.DataFrame,
         embedding_series: pd.Series,
     ) -> "NaturalnessZeroShotModel":
+        if self.is_pretrained:
+            raise ValueError("Model is already pretrained.")
 
         self.knns = {}
         for naturalness_column in naturalness_df.columns:
             naturalness_series = naturalness_df[naturalness_column]
             assert naturalness_series.index.equals(embedding_series.index)
             assert embedding_series.index.is_unique, "embedding_series contains duplicate indices"
 
-            if self.is_pretrained:
-                raise ValueError("Model is already pretrained.")
 
             assert not naturalness_series.isna().any(), "naturalness_series contains NANs"
 
-            logging.info(
-                f"Pretraining model with naturalness data with {len(naturalness_series)} naturalness measurements."
-            )
-
             X = np.array([np.array(emb) for emb in embedding_series.values])
             y = naturalness_series.values
 
             # Fit KNN regressor
             knn = KNeighborsRegressor(n_neighbors=5)
             knn.fit(X, y)
-            self.is_pretrained = True
             self.knns[naturalness_column] = knn
 
         # Also store the naturalness series for prediction later.
-        self.single_mutant_naturalness_ensemble = naturalness_df
+        self.single_mutant_naturalness_df = naturalness_df
+        self.is_pretrained = True
 
         return self
 
@@ -279,9 +262,9 @@ def predict(
         Returns:
             Array of prediction scores based on naturalness
         """
-        assert self.single_mutant_naturalness_ensemble is not None
+        assert self.single_mutant_naturalness_df is not None
         assert len(naturalness_df.columns) == len(self.knns)
-        assert len(naturalness_df.columns) == len(self.single_mutant_naturalness_ensemble)
+        assert set(naturalness_df.columns) == set(self.single_mutant_naturalness_df.columns)
 
         # Get the ensemble of naturalness scores with missing values imputed.
 
@@ -290,7 +273,7 @@ def predict(
         for naturalness_column in naturalness_df.columns:
 
             naturalness_series = naturalness_df[naturalness_column]
-            single_mutant_naturalness_series = self.single_mutant_naturalness_ensemble[
+            single_mutant_naturalness_series = self.single_mutant_naturalness_df[
                 naturalness_column
             ]
             knn = self.knns[naturalness_column]
@@ -314,6 +297,7 @@ def get_naturalness(seq_id, direct_naturalness) -> float:
                     )
                 return computed_naturalness
 
+            naturalness_series.index.name = 'seq_id'
             computed_naturalness_series = naturalness_series.reset_index(name="wt_marginal").apply(
                 lambda r: get_naturalness(r.seq_id, r.wt_marginal), axis=1
             )

backend/src/notebooks/jacob/250713_naturalness_param_scan.py

@@ -64,6 +64,10 @@
         ModelDiff(name="650m", diffs={"naturalness_model_id": "650m"}),
         ModelDiff(name="3b", diffs={"naturalness_model_id": "3b"}),
         ModelDiff(name="15b", diffs={"naturalness_model_id": "15b"}),
+        ModelDiff(name="v1", diffs={
+            "naturalness_model_id": "v1",
+            "few_shot_naturaless_column": ['wt_marginal_1', 'wt_marginal_2', 'wt_marginal_3', 'wt_marginal_4', 'wt_marginal_5'],
+        }),
     ],
 )
 

backend/src/notebooks/jacob/250718_train_ablation.py

@@ -196,6 +196,17 @@
                 "few_shot_model_params.shrink_and_perturb_params": [0.8, 0.0001],
             },
         ),
+        ModelDiff(name="with-esmv1", diffs={
+            "naturalness_model_id": "v1",
+            "few_shot_naturaless_column": ['wt_marginal_1', 'wt_marginal_2', 'wt_marginal_3', 'wt_marginal_4', 'wt_marginal_5'],
+        }),
+        ModelDiff(name="with-esmv1-zsclS025", diffs={
+            "naturalness_model_id": "v1",
+            "few_shot_naturaless_column": ['wt_marginal_1', 'wt_marginal_2', 'wt_marginal_3', 'wt_marginal_4', 'wt_marginal_5'],
+
+            "zero_shot_model_params.decision_mode": "constantliar",
+            "zero_shot_model_params.lie_noise_stddev_multiplier": 0.25,
+        }),
     ],
 )
 

backend/src/notebooks/jacob/250917_zero_shot_cl.py

@@ -0,0 +1,126 @@
+# Configure logging
+import logging
+
+logging.basicConfig(
+    level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+)
+logger = logging.getLogger()
+
+from folde.campaign import simulate_campaigns_with_config_checkpoints
+from folde.types import FolDEModelConfig, ModelDiff
+from folde.util import apply_diff_list_to_config
+
+folde_train_dms_ids = [
+    "ANCSZ_Hobbs_2022",
+    "BLAT_ECOLX_Firnberg_2014",
+    "CBS_HUMAN_Sun_2020",
+    "HEM3_HUMAN_Loggerenberg_2023",
+    "HSP82_YEAST_Flynn_2019",
+    "HXK4_HUMAN_Gersing_2022_activity",
+    "OXDA_RHOTO_Vanella_2023_activity",
+    "PPM1D_HUMAN_Miller_2022",
+    "SHOC2_HUMAN_Kwon_2022",
+]
+
+# Example configuration
+NAME = "250917_zero_shot_cl"
+
+random_config = FolDEModelConfig(
+    name="Random",
+    naturalness_model_id="600m",
+    embedding_model_id="300m",
+    zero_shot_model_name="RandomZeroShotModel",
+    zero_shot_model_params={},
+    few_shot_model_name="RandomFewShotModel",
+    few_shot_model_params={},
+)
+
+random_forest_config = FolDEModelConfig(
+    name="RandomToRandomForest",
+    naturalness_model_id="600m",
+    embedding_model_id="300m",
+    zero_shot_model_name="RandomZeroShotModel",
+    zero_shot_model_params={},
+    few_shot_model_name="RandomForestFewShotModel",
+    few_shot_model_params={
+        "n_estimators": 100,
+        "criterion": "friedman_mse",
+        "max_depth": None,
+        "min_samples_split": 2,
+        "min_samples_leaf": 1,
+        "min_weight_fraction_leaf": 0.0,
+        "max_features": 1.0,
+        "max_leaf_nodes": None,
+        "min_impurity_decrease": 0.0,
+        "bootstrap": True,
+        "oob_score": False,
+        "n_jobs": None,
+        "verbose": 0,
+        "warm_start": False,
+        "ccp_alpha": 0.0,
+        "max_samples": None,
+    },
+)
+
+folde_config = FolDEModelConfig(
+    name="FolDE",
+    # Required parameters
+    naturalness_model_id="600m",  # ESM-2 650M model
+    embedding_model_id="300m",  # Same model for embeddings
+    zero_shot_model_name="NaturalnessZeroShotModel",
+    zero_shot_model_params={},
+    # Few-shot model configuration (used after first round)
+    few_shot_model_name="TorchMLPFewShotModel",
+    few_shot_model_params={
+        "pretrain": True,
+        "pretrain_epochs": 50,
+        "ensemble_size": 5,
+        "embedding_dim": 960,
+        "hidden_dims": [100, 50],
+        "dropout": 0.2,
+        "learning_rate": 3e-4,
+        "weight_decay": 1e-5,
+        "train_epochs": 200,
+        "train_patience": 40,
+        "val_frequency": 10,
+        "do_validation_with_pair_fraction": 0.2,
+        "decision_mode": "constantliar",
+        "lie_noise_stddev_multiplier": 0.25,
+    },
+)
+
+mlp_config_list = apply_diff_list_to_config(
+    folde_config,
+    [
+        ModelDiff(name="no-constantliar", diffs={"few_shot_model_params.decision_mode": "mean"}),
+        ModelDiff(name="with-esmv1", diffs={
+            "naturalness_model_id": "1v",
+            "naturalness_columns": ['wt_marginal_1', 'wt_marginal_2', 'wt_marginal_3', 'wt_marginal_4', 'wt_marginal_5'],
+        }),
+        ModelDiff(name="with-esmv1-zsclS025", diffs={
+            "naturalness_model_id": "1v",
+            "naturalness_columns": ['wt_marginal_1', 'wt_marginal_2', 'wt_marginal_3', 'wt_marginal_4', 'wt_marginal_5'],
+
+            "zero_shot_model_params.decision_mode": "constantliar",
+            "zero_shot_model_params.lie_noise_stddev_multiplier": 0.25,
+        }),
+    ],
+)
+
+config_list = [random_config, random_forest_config] + mlp_config_list
+
+print(f"Config 1/{len(config_list)}:")
+print(config_list[0].model_dump_json(indent=2))
+
+results = simulate_campaigns_with_config_checkpoints(
+    eval_prefix=NAME,
+    dms_ids=folde_train_dms_ids,
+    config_list=config_list,
+    checkpoint_dir="notebooks/jacob/model_evals",
+    round_size=16,
+    number_of_simulations=10,
+    activity_column="DMS_score",
+    max_rounds=6,
+    random_seed=42,
+    num_workers=10,
+)