add electra freeze option

Author: sfluegel05

chebai/models/electra.py

@@ -224,6 +224,7 @@ def __init__(
         config: Optional[Dict[str, Any]] = None,
         pretrained_checkpoint: Optional[str] = None,
         load_prefix: Optional[str] = None,
+        freeze_electra: bool = False,
         **kwargs: Any,
     ):
         # Remove this property in order to prevent it from being stored as a
@@ -262,6 +263,10 @@ def __init__(
         else:
             self.electra = ElectraModel(config=self.config)
 
+        if freeze_electra:
+            for param in self.electra.parameters():
+                param.requires_grad = False
+
     def _process_for_loss(
         self,
         model_output: Dict[str, Tensor],