| 266 | |
| 267 | class ReprogrammingLayer(nn.Module): |
| 268 | def __init__(self, d_model, n_heads, d_keys=None, d_llm=None, attention_dropout=0.1): |
| 269 | super(ReprogrammingLayer, self).__init__() |
| 270 | |
| 271 | d_keys = d_keys or (d_model // n_heads) |
| 272 | |
| 273 | self.query_projection = nn.Linear(d_model, d_keys * n_heads) |
| 274 | self.key_projection = nn.Linear(d_llm, d_keys * n_heads) |
| 275 | self.value_projection = nn.Linear(d_llm, d_keys * n_heads) |
| 276 | self.out_projection = nn.Linear(d_keys * n_heads, d_llm) |
| 277 | self.n_heads = n_heads |
| 278 | self.dropout = nn.Dropout(attention_dropout) |
| 279 | |
| 280 | def forward(self, target_embedding, source_embedding, value_embedding): |
| 281 | B, L, _ = target_embedding.shape |