add quantization for pretrained model

geneformer/in_silico_perturber.py

@@ -62,7 +62,7 @@ class InSilicoPerturber:
         "genes_to_perturb": {"all", list},
         "combos": {0, 1},
         "anchor_gene": {None, str},
-        "model_type": {"Pretrained", "GeneClassifier", "CellClassifier", "MTLCellClassifier", "MTLCellClassifier-Quantized"},
+        "model_type": {"Pretrained", "GeneClassifier", "CellClassifier", "MTLCellClassifier", "Pretrained-Quantized", "MTLCellClassifier-Quantized"},
         "num_classes": {int},
         "emb_mode": {"cls", "cell", "cls_and_gene", "cell_and_gene"},
         "cell_emb_style": {"mean_pool"},
@@ -132,7 +132,7 @@ class InSilicoPerturber:
             | ENSEMBL ID of gene to use as anchor in combination perturbations.
             | For example, if combos=1 and anchor_gene="ENSG00000148400":
             |     anchor gene will be perturbed in combination with each other gene.
-        model_type : {"Pretrained", "GeneClassifier", "CellClassifier", "MTLCellClassifier", "MTLCellClassifier-Quantized"}
+        model_type : {"Pretrained", "GeneClassifier", "CellClassifier", "MTLCellClassifier", "Pretrained-Quantized", "MTLCellClassifier-Quantized"}
             | Whether model is the pretrained Geneformer or a fine-tuned gene, cell, or multitask cell classifier (+/- 8bit quantization).
         num_classes : int
             | If model is a gene or cell classifier, specify number of classes it was trained to classify.

geneformer/perturber_utils.py

@@ -113,15 +113,22 @@ def slice_by_inds_to_perturb(filtered_input_data, cell_inds_to_perturb):
 
 # load model to GPU
 def load_model(model_type, num_classes, model_directory, mode, quantize=False):
-    if model_type == "MTLCellClassifier-Quantized":
+    if model_type == "Pretrained-Quantized":
+        inference_only = True
+        model_type = "Pretrained"
+        quantize = True
+    elif model_type == "MTLCellClassifier-Quantized":
+        inference_only = True
         model_type = "MTLCellClassifier"
         quantize = True
+    else:
+        inference_only = False
 
     output_hidden_states = (mode == "eval")
 
     # Quantization logic
     if quantize:
-        if model_type == "MTLCellClassifier":
+        if inference_only:
             quantize_config = BitsAndBytesConfig(load_in_8bit=True)
             peft_config = None
         else:
@@ -179,7 +186,7 @@ def load_model(model_type, num_classes, model_directory, mode, quantize=False):
         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         model = model.to(device)
     elif peft_config:
-        # Apply PEFT for quantized models (except MTLCellClassifier)
+        # Apply PEFT for quantized models (except MTLCellClassifier and CellClassifier-QuantInf)
         model.enable_input_require_grads()
         model = get_peft_model(model, peft_config)