The torch.nn model to encode the prefix Input shape: (batch-size, prefix-length) Output shape: (batch-size, prefix-length, 2*layers*hidden)
| 59 | |
| 60 | |
| 61 | class PrefixEncoder(torch.nn.Module): |
| 62 | """ |
| 63 | The torch.nn model to encode the prefix |
| 64 | Input shape: (batch-size, prefix-length) |
| 65 | Output shape: (batch-size, prefix-length, 2*layers*hidden) |
| 66 | """ |
| 67 | |
| 68 | def __init__(self, config: ChatGLMConfig): |
| 69 | super().__init__() |
| 70 | self.prefix_projection = config.prefix_projection |
| 71 | if self.prefix_projection: |
| 72 | # Use a two-layer MLP to encode the prefix |
| 73 | kv_size = config.num_layers * config.kv_channels * config.multi_query_group_num * 2 |
| 74 | self.embedding = torch.nn.Embedding(config.pre_seq_len, kv_size) |
| 75 | self.trans = torch.nn.Sequential( |
| 76 | torch.nn.Linear(kv_size, config.hidden_size), |
| 77 | torch.nn.Tanh(), |
| 78 | torch.nn.Linear(config.hidden_size, kv_size) |
| 79 | ) |
| 80 | else: |
| 81 | self.embedding = torch.nn.Embedding(config.pre_seq_len, |
| 82 | config.num_layers * config.kv_channels * config.multi_query_group_num * 2) |
| 83 | |
| 84 | def forward(self, prefix: torch.Tensor): |
| 85 | if self.prefix_projection: |
| 86 | prefix_tokens = self.embedding(prefix) |
| 87 | past_key_values = self.trans(prefix_tokens) |
| 88 | else: |
| 89 | past_key_values = self.embedding(prefix) |
| 90 | return past_key_values |
| 91 | |
| 92 | |
| 93 | def split_tensor_along_last_dim( |