| 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: |