| 324 | """ |
| 325 | |
| 326 | def __init__(self, config: ChatGLMConfig): |
| 327 | super().__init__() |
| 328 | self.prefix_projection = config.prefix_projection |
| 329 | if self.prefix_projection: |
| 330 | # Use a two-layer MLP to encode the prefix |
| 331 | kv_size = config.num_layers * config.kv_channels * config.multi_query_group_num * 2 |
| 332 | self.embedding = torch.nn.Embedding(config.pre_seq_len, kv_size) |
| 333 | self.trans = torch.nn.Sequential( |
| 334 | torch.nn.Linear(kv_size, config.hidden_size), |
| 335 | torch.nn.Tanh(), |
| 336 | torch.nn.Linear(config.hidden_size, kv_size) |
| 337 | ) |
| 338 | else: |
| 339 | self.embedding = torch.nn.Embedding(config.pre_seq_len, |
| 340 | config.num_layers * config.kv_channels * config.multi_query_group_num * 2) |
| 341 | |
| 342 | def forward(self, prefix: torch.Tensor): |
| 343 | if self.prefix_projection: |