| from transformers import PretrainedConfig |
|
|
|
|
| class SingProbeMlpConfig(PretrainedConfig): |
| model_type = "sing_probe_mlp" |
|
|
| def __init__( |
| self, |
| hidden_size: int = 2560, |
| base_model_layer_ids: list[int] | None = None, |
| intermediate_size: int = 1024, |
| num_labels: int = 10, |
| hidden_act: str = "gelu", |
| base_model_name: str | None = None, |
| **kwargs, |
| ) -> None: |
| super().__init__(**kwargs) |
| self.hidden_size = int(hidden_size) |
| self.base_model_layer_ids = base_model_layer_ids or [] |
| self.intermediate_size = int(intermediate_size) |
| self.num_labels = int(num_labels) |
| self.hidden_act = hidden_act |
| self.base_model_name = base_model_name |
|
|
| @property |
| def input_size(self) -> int: |
| return self.hidden_size * len(self.base_model_layer_ids) |
|
|
|
|
| class SingProbeAttnConfig(PretrainedConfig): |
| model_type = "sing_probe_attn" |
|
|
| def __init__( |
| self, |
| hidden_size: int = 2560, |
| base_model_layer_ids: list[int] | None = None, |
| num_attention_heads: int = 4, |
| head_dim: int = 64, |
| sliding_window: int | None = None, |
| num_labels: int = 10, |
| base_model_name: str | None = None, |
| **kwargs, |
| ) -> None: |
| super().__init__(**kwargs) |
| self.hidden_size = int(hidden_size) |
| self.base_model_layer_ids = base_model_layer_ids or [] |
| self.num_attention_heads = int(num_attention_heads) |
| self.head_dim = int(head_dim) |
| self.sliding_window = None if sliding_window is None else int(sliding_window) |
| self.num_labels = int(num_labels) |
| self.base_model_name = base_model_name |
|
|
| @property |
| def input_size(self) -> int: |
| return self.hidden_size * len(self.base_model_layer_ids) |
|
|
