File size: 1,807 Bytes
4bb52fb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 | 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)
|