The torch.nn model to encode the prefix Input shape: (batch-size, prefix-length) Output shape: (batch-size, prefix-length, 2*layers*hidden)
| 317 | |
| 318 | |
| 319 | class PrefixEncoder(torch.nn.Module): |
| 320 | """ |
| 321 | The torch.nn model to encode the prefix |
| 322 | Input shape: (batch-size, prefix-length) |
| 323 | Output shape: (batch-size, prefix-length, 2*layers*hidden) |
| 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: |
| 344 | prefix_tokens = self.embedding(prefix) |
| 345 | past_key_values = self.trans(prefix_tokens) |
| 346 | else: |
| 347 | past_key_values = self.embedding(prefix) |
| 348 | return past_key_values |
| 349 | |
| 350 | |
| 351 | def split_tensor_along_last_dim( |